Skip to content

Commit 043db8b

Browse files
njacazionjacazio
authored
Track propagation with ID (#3989)
* Track propagation with ID * Update trk utils
1 parent e6664be commit 043db8b

2 files changed

Lines changed: 156 additions & 103 deletions

File tree

‎Common/Core/trackUtilities.h‎

Lines changed: 37 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -25,26 +25,50 @@
2525
template <typename TrackPrecision = float, typename T>
2626
o2::track::TrackParametrization<TrackPrecision> getTrackPar(const T& track)
2727
{
28-
std::array<TrackPrecision, 5> arraypar = {track.y(), track.z(), track.snp(),
29-
track.tgl(), track.signed1Pt()};
28+
std::array<TrackPrecision, o2::track::kNParams> arraypar = {track.y(), track.z(), track.snp(),
29+
track.tgl(), track.signed1Pt()};
3030
return o2::track::TrackParametrization<TrackPrecision>(track.x(), track.alpha(), std::move(arraypar));
3131
}
3232

33+
/// Extracts track parameters from a track and sets a TrackParametrization object.
34+
template <typename TrackPrecision = float, typename T>
35+
void setTrackPar(const T& track, o2::track::TrackParametrization<TrackPrecision>& trackPar)
36+
{
37+
std::array<TrackPrecision, o2::track::kNParams> arraypar = {track.y(), track.z(), track.snp(),
38+
track.tgl(), track.signed1Pt()};
39+
trackPar.set(track.x(), track.alpha(), std::move(arraypar));
40+
}
41+
3342
/// Extracts track parameters and covariance matrix from a track.
3443
template <typename TrackPrecision = float, typename T>
3544
o2::track::TrackParametrizationWithError<TrackPrecision> getTrackParCov(const T& track)
3645
{
37-
std::array<TrackPrecision, 5> arraypar = {track.y(), track.z(), track.snp(),
38-
track.tgl(), track.signed1Pt()};
39-
std::array<TrackPrecision, 15> covpar = {track.cYY(), track.cZY(), track.cZZ(),
40-
track.cSnpY(), track.cSnpZ(),
41-
track.cSnpSnp(), track.cTglY(), track.cTglZ(),
42-
track.cTglSnp(), track.cTglTgl(),
43-
track.c1PtY(), track.c1PtZ(), track.c1PtSnp(),
44-
track.c1PtTgl(), track.c1Pt21Pt2()};
46+
std::array<TrackPrecision, o2::track::kNParams> arraypar = {track.y(), track.z(), track.snp(),
47+
track.tgl(), track.signed1Pt()};
48+
std::array<TrackPrecision, o2::track::kCovMatSize> covpar = {track.cYY(), track.cZY(), track.cZZ(),
49+
track.cSnpY(), track.cSnpZ(),
50+
track.cSnpSnp(), track.cTglY(), track.cTglZ(),
51+
track.cTglSnp(), track.cTglTgl(),
52+
track.c1PtY(), track.c1PtZ(), track.c1PtSnp(),
53+
track.c1PtTgl(), track.c1Pt21Pt2()};
4554
return o2::track::TrackParametrizationWithError<TrackPrecision>(track.x(), track.alpha(), std::move(arraypar), std::move(covpar));
4655
}
4756

57+
/// Extracts track parameters and covariance matrix from a track and sets a TrackParametrizationWithError object.
58+
template <typename TrackPrecision = float, typename T>
59+
void setTrackParCov(const T& track, o2::track::TrackParametrizationWithError<TrackPrecision>& trackParCov)
60+
{
61+
std::array<TrackPrecision, o2::track::kNParams> arraypar = {track.y(), track.z(), track.snp(),
62+
track.tgl(), track.signed1Pt()};
63+
std::array<TrackPrecision, o2::track::kCovMatSize> covpar = {track.cYY(), track.cZY(), track.cZZ(),
64+
track.cSnpY(), track.cSnpZ(),
65+
track.cSnpSnp(), track.cTglY(), track.cTglZ(),
66+
track.cTglSnp(), track.cTglTgl(),
67+
track.c1PtY(), track.c1PtZ(), track.c1PtSnp(),
68+
track.c1PtTgl(), track.c1Pt21Pt2()};
69+
trackParCov.set(track.x(), track.alpha(), std::move(arraypar), std::move(covpar));
70+
}
71+
4872
/// Extracts primary vertex position and covariance matrix from a collision.
4973
template <typename T>
5074
o2::dataformats::VertexBase getPrimaryVertex(const T& collision)
@@ -63,9 +87,9 @@ template <typename T, typename U, typename V, typename W>
6387
void getPointDirection(const T& point1, const U& point2, V& phi, W& theta)
6488
{
6589
phi = std::atan2(point2[1] - point1[1], point2[0] - point1[0]);
66-
//auto x1 = point1[0] * std::cos(phi) + point1[1] * std::sin(phi);
67-
//auto x2 = point2[0] * std::cos(phi) + point2[1] * std::sin(phi);
68-
//theta = std::atan2(point2[2] - point1[2], x2 - x1);
90+
// auto x1 = point1[0] * std::cos(phi) + point1[1] * std::sin(phi);
91+
// auto x2 = point2[0] * std::cos(phi) + point2[1] * std::sin(phi);
92+
// theta = std::atan2(point2[2] - point1[2], x2 - x1);
6993
theta = std::atan2(point2[2] - point1[2], RecoDecay::distanceXY(point1, point2));
7094
}
7195

‎Common/TableProducer/trackPropagation.cxx‎

Lines changed: 119 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "Framework/HistogramRegistry.h"
3030
#include "DataFormatsCalibration/MeanVertexObject.h"
3131
#include "CommonConstants/GeomConstants.h"
32+
#include "TableHelper.h"
3233

3334
// The Run 3 AO2D stores the tracks at the point of innermost update. For a track with ITS this is the innermost (or second innermost)
3435
// ITS layer. For a track without ITS, this is the TPC inner wall or for loopers in the TPC even a radius beyond that.
@@ -60,7 +61,7 @@ struct TrackPropagation {
6061

6162
o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT;
6263

63-
const o2::dataformats::MeanVertexObject* mVtx = nullptr;
64+
const o2::dataformats::MeanVertexObject* mMeanVtx = nullptr;
6465
o2::parameters::GRPMagField* grpmag = nullptr;
6566
o2::base::MatLayerCylSet* lut = nullptr;
6667

@@ -73,22 +74,29 @@ struct TrackPropagation {
7374

7475
void init(o2::framework::InitContext& initContext)
7576
{
76-
if (doprocessCovariance == true && doprocessStandard == true) {
77-
LOGF(fatal, "Cannot enable processStandard and processCovariance at the same time. Please choose one.");
77+
int nEnabledProcesses = 0;
78+
if (doprocessStandard == true) {
79+
LOG(info) << "Enabling processStandard";
80+
nEnabledProcesses++;
7881
}
79-
80-
// Checking if the tables are requested in the workflow and enabling them
81-
auto& workflows = initContext.services().get<RunningWorkflowInfo const>();
82-
for (DeviceSpec const& device : workflows.devices) {
83-
for (auto const& input : device.inputs) {
84-
if (input.matcher.binding == "TracksDCA") {
85-
fillTracksDCA = true;
86-
}
87-
if (input.matcher.binding == "TracksDCACov") {
88-
fillTracksDCACov = true;
89-
}
90-
}
82+
if (doprocessCovariance == true) {
83+
LOG(info) << "Enabling processCovariance";
84+
nEnabledProcesses++;
85+
}
86+
if (doprocessStandardWithPID == true) {
87+
LOG(info) << "Enabling processStandardWithPID";
88+
nEnabledProcesses++;
89+
}
90+
if (doprocessCovarianceWithPID == true) {
91+
LOG(info) << "Enabling processCovarianceWithPID";
92+
nEnabledProcesses++;
93+
}
94+
if (nEnabledProcesses != 1) {
95+
LOG(fatal) << "Exactly one process flag must be set to true. Please choose one.";
9196
}
97+
// Checking if the tables are requested in the workflow and enabling them
98+
fillTracksDCA = isTableRequiredInWorkflow(initContext, "TracksDCA");
99+
fillTracksDCACov = isTableRequiredInWorkflow(initContext, "TracksDCACov");
92100

93101
ccdb->setURL(ccdburl);
94102
ccdb->setCaching(true);
@@ -106,111 +114,132 @@ struct TrackPropagation {
106114
LOG(info) << "Setting magnetic field to current " << grpmag->getL3Current() << " A for run " << bc.runNumber() << " from its GRPMagField CCDB object";
107115
o2::base::Propagator::initFieldFromGRP(grpmag);
108116
o2::base::Propagator::Instance()->setMatLUT(lut);
109-
mVtx = ccdb->getForTimeStamp<o2::dataformats::MeanVertexObject>(mVtxPath, bc.timestamp());
117+
mMeanVtx = ccdb->getForTimeStamp<o2::dataformats::MeanVertexObject>(mVtxPath, bc.timestamp());
110118
runNumber = bc.runNumber();
111119
}
112120

113-
template <typename TTrack, typename TTrackPar>
114-
void FillTracksPar(TTrack& track, aod::track::TrackTypeEnum trackType, TTrackPar& trackPar)
115-
{
116-
tracksParPropagated(track.collisionId(), trackType, trackPar.getX(), trackPar.getAlpha(), trackPar.getY(), trackPar.getZ(), trackPar.getSnp(), trackPar.getTgl(), trackPar.getQ2Pt());
117-
tracksParExtensionPropagated(trackPar.getPt(), trackPar.getP(), trackPar.getEta(), trackPar.getPhi());
118-
}
119-
120-
void processStandard(aod::StoredTracksIU const& tracks, aod::Collisions const&, aod::BCsWithTimestamps const& bcs)
121+
// Running variables
122+
gpu::gpustd::array<float, 2> mDcaInfo;
123+
o2::dataformats::DCA mDcaInfoCov;
124+
o2::dataformats::VertexBase mVtx;
125+
o2::track::TrackParametrization<float> mTrackPar;
126+
o2::track::TrackParametrizationWithError<float> mTrackParCov;
127+
128+
template <typename TTrack, bool fillCovMat = false, bool useTrkPid = false>
129+
void fillTrackTables(TTrack const& tracks,
130+
aod::Collisions const&,
131+
aod::BCsWithTimestamps const& bcs)
121132
{
122133
if (bcs.size() == 0) {
123134
return;
124135
}
125136
initCCDB(bcs.begin());
126137

127-
tracksParPropagated.reserve(tracks.size());
128-
tracksParExtensionPropagated.reserve(tracks.size());
129-
if (fillTracksDCA) {
130-
tracksDCA.reserve(tracks.size());
138+
if constexpr (fillCovMat) {
139+
tracksParCovPropagated.reserve(tracks.size());
140+
tracksParCovExtensionPropagated.reserve(tracks.size());
141+
if (fillTracksDCACov) {
142+
tracksDCACov.reserve(tracks.size());
143+
}
144+
} else {
145+
tracksParPropagated.reserve(tracks.size());
146+
tracksParExtensionPropagated.reserve(tracks.size());
147+
if (fillTracksDCA) {
148+
tracksDCA.reserve(tracks.size());
149+
}
131150
}
132151

133-
gpu::gpustd::array<float, 2> dcaInfo;
134-
135152
for (auto& track : tracks) {
136-
dcaInfo[0] = 999;
137-
dcaInfo[1] = 999;
153+
if constexpr (fillCovMat) {
154+
if (fillTracksDCACov) {
155+
mDcaInfoCov.set(999, 999, 999, 999, 999);
156+
}
157+
setTrackParCov(track, mTrackParCov);
158+
if constexpr (useTrkPid) {
159+
mTrackParCov.setPID(track.pidForTracking());
160+
}
161+
} else {
162+
if (fillTracksDCA) {
163+
mDcaInfo[0] = 999;
164+
mDcaInfo[1] = 999;
165+
}
166+
setTrackPar(track, mTrackPar);
167+
if constexpr (useTrkPid) {
168+
mTrackPar.setPID(track.pidForTracking());
169+
}
170+
}
138171
aod::track::TrackTypeEnum trackType = (aod::track::TrackTypeEnum)track.trackType();
139-
auto trackPar = getTrackPar(track);
140172
// Only propagate tracks which have passed the innermost wall of the TPC (e.g. skipping loopers etc). Others fill unpropagated.
141173
if (track.trackType() == aod::track::TrackIU && track.x() < minPropagationRadius) {
142174
if (track.has_collision()) {
143175
auto const& collision = track.collision();
144-
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackPar, 2.f, matCorr, &dcaInfo);
176+
if constexpr (fillCovMat) {
177+
mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()});
178+
mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ());
179+
o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, mTrackParCov, 2.f, matCorr, &mDcaInfoCov);
180+
} else {
181+
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, mTrackPar, 2.f, matCorr, &mDcaInfo);
182+
}
145183
} else {
146-
o2::base::Propagator::Instance()->propagateToDCABxByBz({mVtx->getX(), mVtx->getY(), mVtx->getZ()}, trackPar, 2.f, matCorr, &dcaInfo);
184+
if constexpr (fillCovMat) {
185+
mVtx.setPos({mMeanVtx->getX(), mMeanVtx->getY(), mMeanVtx->getZ()});
186+
mVtx.setCov(mMeanVtx->getSigmaX() * mMeanVtx->getSigmaX(), 0.0f, mMeanVtx->getSigmaY() * mMeanVtx->getSigmaY(), 0.0f, 0.0f, mMeanVtx->getSigmaZ() * mMeanVtx->getSigmaZ());
187+
o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, mTrackParCov, 2.f, matCorr, &mDcaInfoCov);
188+
} else {
189+
o2::base::Propagator::Instance()->propagateToDCABxByBz({mMeanVtx->getX(), mMeanVtx->getY(), mMeanVtx->getZ()}, mTrackPar, 2.f, matCorr, &mDcaInfo);
190+
}
147191
}
148192
trackType = aod::track::Track;
149193
}
150-
FillTracksPar(track, trackType, trackPar);
151-
if (fillTracksDCA) {
152-
tracksDCA(dcaInfo[0], dcaInfo[1]);
194+
if constexpr (fillCovMat) {
195+
tracksParPropagated(track.collisionId(), trackType, mTrackParCov.getX(), mTrackParCov.getAlpha(), mTrackParCov.getY(), mTrackParCov.getZ(), mTrackParCov.getSnp(), mTrackParCov.getTgl(), mTrackParCov.getQ2Pt());
196+
tracksParExtensionPropagated(mTrackParCov.getPt(), mTrackParCov.getP(), mTrackParCov.getEta(), mTrackParCov.getPhi());
197+
// TODO do we keep the rho as 0? Also the sigma's are duplicated information
198+
tracksParCovPropagated(std::sqrt(mTrackParCov.getSigmaY2()), std::sqrt(mTrackParCov.getSigmaZ2()), std::sqrt(mTrackParCov.getSigmaSnp2()),
199+
std::sqrt(mTrackParCov.getSigmaTgl2()), std::sqrt(mTrackParCov.getSigma1Pt2()), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
200+
tracksParCovExtensionPropagated(mTrackParCov.getSigmaY2(), mTrackParCov.getSigmaZY(), mTrackParCov.getSigmaZ2(), mTrackParCov.getSigmaSnpY(),
201+
mTrackParCov.getSigmaSnpZ(), mTrackParCov.getSigmaSnp2(), mTrackParCov.getSigmaTglY(), mTrackParCov.getSigmaTglZ(), mTrackParCov.getSigmaTglSnp(),
202+
mTrackParCov.getSigmaTgl2(), mTrackParCov.getSigma1PtY(), mTrackParCov.getSigma1PtZ(), mTrackParCov.getSigma1PtSnp(), mTrackParCov.getSigma1PtTgl(),
203+
mTrackParCov.getSigma1Pt2());
204+
if (fillTracksDCA) {
205+
tracksDCA(mDcaInfoCov.getY(), mDcaInfoCov.getZ());
206+
}
207+
if (fillTracksDCACov) {
208+
tracksDCACov(mDcaInfoCov.getSigmaY2(), mDcaInfoCov.getSigmaZ2());
209+
}
210+
} else {
211+
tracksParPropagated(track.collisionId(), trackType, mTrackPar.getX(), mTrackPar.getAlpha(), mTrackPar.getY(), mTrackPar.getZ(), mTrackPar.getSnp(), mTrackPar.getTgl(), mTrackPar.getQ2Pt());
212+
tracksParExtensionPropagated(mTrackPar.getPt(), mTrackPar.getP(), mTrackPar.getEta(), mTrackPar.getPhi());
213+
if (fillTracksDCA) {
214+
tracksDCA(mDcaInfo[0], mDcaInfo[1]);
215+
}
153216
}
154217
}
155218
}
156-
PROCESS_SWITCH(TrackPropagation, processStandard, "Process without covariance", true);
157219

158-
void processCovariance(soa::Join<aod::StoredTracksIU, aod::TracksCovIU> const& tracks, aod::Collisions const&, aod::BCsWithTimestamps const& bcs)
220+
void processStandard(aod::StoredTracksIU const& tracks, aod::Collisions const& collisions, aod::BCsWithTimestamps const& bcs)
159221
{
160-
if (bcs.size() == 0) {
161-
return;
162-
}
163-
initCCDB(bcs.begin());
164-
165-
o2::dataformats::DCA dcaInfoCov;
166-
o2::dataformats::VertexBase vtx;
222+
fillTrackTables<aod::StoredTracksIU, /*fillCovMat =*/false, /*useTrkPid =*/false>(tracks, collisions, bcs);
223+
}
224+
PROCESS_SWITCH(TrackPropagation, processStandard, "Process without covariance", true);
167225

168-
tracksParPropagated.reserve(tracks.size());
169-
tracksParExtensionPropagated.reserve(tracks.size());
170-
tracksParCovPropagated.reserve(tracks.size());
171-
tracksParCovExtensionPropagated.reserve(tracks.size());
172-
if (fillTracksDCA) {
173-
tracksDCA.reserve(tracks.size());
174-
}
175-
if (fillTracksDCACov) {
176-
tracksDCACov.reserve(tracks.size());
177-
}
226+
void processStandardWithPID(soa::Join<aod::StoredTracksIU, aod::TracksExtra> const& tracks, aod::Collisions const& collisions, aod::BCsWithTimestamps const& bcs)
227+
{
228+
fillTrackTables<soa::Join<aod::StoredTracksIU, aod::TracksExtra>, /*fillCovMat =*/false, /*useTrkPid =*/true>(tracks, collisions, bcs);
229+
}
230+
PROCESS_SWITCH(TrackPropagation, processStandardWithPID, "Process without covariance and with PID in tracking", false);
178231

179-
for (auto& track : tracks) {
180-
dcaInfoCov.set(999, 999, 999, 999, 999);
181-
auto trackParCov = getTrackParCov(track);
182-
aod::track::TrackTypeEnum trackType = (aod::track::TrackTypeEnum)track.trackType();
183-
// Only propagate tracks which have passed the innermost wall of the TPC (e.g. skipping loopers etc). Others fill unpropagated.
184-
if (track.trackType() == aod::track::TrackIU && track.x() < minPropagationRadius) {
185-
if (track.has_collision()) {
186-
auto const& collision = track.collision();
187-
vtx.setPos({collision.posX(), collision.posY(), collision.posZ()});
188-
vtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ());
189-
o2::base::Propagator::Instance()->propagateToDCABxByBz(vtx, trackParCov, 2.f, matCorr, &dcaInfoCov);
190-
} else {
191-
vtx.setPos({mVtx->getX(), mVtx->getY(), mVtx->getZ()});
192-
vtx.setCov(mVtx->getSigmaX() * mVtx->getSigmaX(), 0.0f, mVtx->getSigmaY() * mVtx->getSigmaY(), 0.0f, 0.0f, mVtx->getSigmaZ() * mVtx->getSigmaZ());
193-
o2::base::Propagator::Instance()->propagateToDCABxByBz(vtx, trackParCov, 2.f, matCorr, &dcaInfoCov);
194-
}
195-
trackType = aod::track::Track;
196-
}
197-
FillTracksPar(track, trackType, trackParCov);
198-
if (fillTracksDCA) {
199-
tracksDCA(dcaInfoCov.getY(), dcaInfoCov.getZ());
200-
}
201-
if (fillTracksDCACov) {
202-
tracksDCACov(dcaInfoCov.getSigmaY2(), dcaInfoCov.getSigmaZ2());
203-
}
204-
// TODO do we keep the rho as 0? Also the sigma's are duplicated information
205-
tracksParCovPropagated(std::sqrt(trackParCov.getSigmaY2()), std::sqrt(trackParCov.getSigmaZ2()), std::sqrt(trackParCov.getSigmaSnp2()),
206-
std::sqrt(trackParCov.getSigmaTgl2()), std::sqrt(trackParCov.getSigma1Pt2()), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
207-
tracksParCovExtensionPropagated(trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(), trackParCov.getSigmaSnpY(),
208-
trackParCov.getSigmaSnpZ(), trackParCov.getSigmaSnp2(), trackParCov.getSigmaTglY(), trackParCov.getSigmaTglZ(), trackParCov.getSigmaTglSnp(),
209-
trackParCov.getSigmaTgl2(), trackParCov.getSigma1PtY(), trackParCov.getSigma1PtZ(), trackParCov.getSigma1PtSnp(), trackParCov.getSigma1PtTgl(),
210-
trackParCov.getSigma1Pt2());
211-
}
232+
void processCovariance(soa::Join<aod::StoredTracksIU, aod::TracksCovIU> const& tracks, aod::Collisions const& collisions, aod::BCsWithTimestamps const& bcs)
233+
{
234+
fillTrackTables<soa::Join<aod::StoredTracksIU, aod::TracksCovIU>, /*fillCovMat =*/true, /*useTrkPid =*/false>(tracks, collisions, bcs);
212235
}
213236
PROCESS_SWITCH(TrackPropagation, processCovariance, "Process with covariance", false);
237+
238+
void processCovarianceWithPID(soa::Join<aod::StoredTracksIU, aod::TracksCovIU, aod::TracksExtra> const& tracks, aod::Collisions const& collisions, aod::BCsWithTimestamps const& bcs)
239+
{
240+
fillTrackTables<soa::Join<aod::StoredTracksIU, aod::TracksCovIU, aod::TracksExtra>, /*fillCovMat =*/true, /*useTrkPid =*/false>(tracks, collisions, bcs);
241+
}
242+
PROCESS_SWITCH(TrackPropagation, processCovarianceWithPID, "Process with covariance and with PID in tracking", false);
214243
};
215244

216245
//****************************************************************************************

0 commit comments

Comments
 (0)