Skip to content

Commit

Permalink
implementation of ambiguity resolution solver from ACTS (#1383)
Browse files Browse the repository at this point in the history
### Briefly, what does this PR introduce?
* implementation of ambiguity resolution solver from ACTS; removing
duplicate tracks from realistic seeding
* new factory called after CKFtracking; taking CKFtracking outputs as
inputs; provides a set of output collections

### What kind of change does this PR introduce?
- [ ] Bug fix (issue #__)
- [x] New feature (issue #__)
- [ ] Documentation update
- [ ] Other: __

### Please check if this PR fulfills the following:
- [ ] Tests for the changes have been added
- [ ] Documentation has been added / updated
- [x] Changes have been communicated to collaborators

### Does this PR introduce breaking changes? What changes might users
need to make to their code?
* No change required; 

### Does this PR change default behavior?
* adding 3 PodIO output collection

---------

Co-authored-by: Minjung Kim <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Minjung Kim <[email protected]>
Co-authored-by: Wouter Deconinck <[email protected]>
Co-authored-by: Dmitry Kalinkin <[email protected]>
Co-authored-by: Minjung Kim <[email protected]>
Co-authored-by: Barak Schmookler <[email protected]>
  • Loading branch information
8 people authored Jun 4, 2024
1 parent 2e074ae commit 1413111
Show file tree
Hide file tree
Showing 7 changed files with 304 additions and 12 deletions.
5 changes: 1 addition & 4 deletions src/algorithms/tracking/ActsToTracks.cc
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

#include <Acts/Definitions/TrackParametrization.hpp>
#include <Acts/Definitions/Units.hpp>
#include <Acts/EventData/MultiTrajectory.hpp>
#include <Acts/EventData/MultiTrajectoryHelpers.hpp>
#include <Acts/EventData/ParticleHypothesis.hpp>
#include <Acts/EventData/TrackStateType.hpp>
Expand All @@ -13,7 +12,6 @@
#include <edm4hep/Vector3f.h>
#include <fmt/core.h>
#include <Eigen/Core>
#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
Expand Down Expand Up @@ -58,8 +56,7 @@ void ActsToTracks::process(const Input& input, const Output& output) const {
}

// Loop over all trajectories in a multiTrajectory
// FIXME: we only retain the first trackTips entry
for (auto trackTip : decltype(trackTips){trackTips.front()}) {
for (auto trackTip : trackTips) {
// Collect the trajectory summary info
auto trajectoryState =
Acts::MultiTrajectoryHelpers::trajectoryState(mj, trackTip);
Expand Down
121 changes: 121 additions & 0 deletions src/algorithms/tracking/AmbiguitySolver.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2024 Minjung Kim, Barak Schmookler
#include "AmbiguitySolver.h"

#include <Acts/AmbiguityResolution/GreedyAmbiguityResolution.hpp>
#include <Acts/EventData/GenericBoundTrackParameters.hpp>
#include <Acts/EventData/MultiTrajectory.hpp>
#include <Acts/EventData/ParticleHypothesis.hpp>
#include <Acts/EventData/SourceLink.hpp>
#include <Acts/EventData/TrackContainer.hpp>
#include <Acts/EventData/TrackProxy.hpp>
#include <Acts/EventData/TrackStatePropMask.hpp>
#include <Acts/EventData/VectorMultiTrajectory.hpp>
#include <Acts/EventData/VectorTrackContainer.hpp>
#include <Acts/Surfaces/Surface.hpp>
#include <ActsExamples/EventData/IndexSourceLink.hpp>
#include <ActsExamples/EventData/Track.hpp>
#include <ActsExamples/EventData/Trajectories.hpp>
#include <boost/container/flat_set.hpp>
#include <boost/container/vector.hpp>
#include <edm4eic/Measurement2DCollection.h>
#include <Eigen/Core>
#include <cstddef>
#include <functional>
#include <utility>

#include "Acts/Utilities/Logger.hpp"
#include "AmbiguitySolverConfig.h"
#include "extensions/spdlog/SpdlogFormatters.h" // IWYU pragma: keep
#include "extensions/spdlog/SpdlogToActs.h"

namespace eicrecon {

Acts::GreedyAmbiguityResolution::Config
transformConfig(const eicrecon::AmbiguitySolverConfig& cfg) {
Acts::GreedyAmbiguityResolution::Config result;
result.maximumSharedHits = cfg.maximum_shared_hits;
result.maximumIterations = cfg.maximum_iterations;
result.nMeasurementsMin = cfg.n_measurements_min;
return result;
}

static std::size_t sourceLinkHash(const Acts::SourceLink& a) {
return static_cast<std::size_t>(a.get<ActsExamples::IndexSourceLink>().index());
}

static bool sourceLinkEquality(const Acts::SourceLink& a, const Acts::SourceLink& b) {
return a.get<ActsExamples::IndexSourceLink>().index() ==
b.get<ActsExamples::IndexSourceLink>().index();
}


AmbiguitySolver::AmbiguitySolver() {}


void AmbiguitySolver::init(std::shared_ptr<spdlog::logger> log) {

m_log = log;
m_acts_logger = eicrecon::getSpdlogLogger("AmbiguitySolver", m_log);
m_acts_cfg = transformConfig(m_cfg);
m_core = std::make_unique<Acts::GreedyAmbiguityResolution>(m_acts_cfg, logger().clone());
}


std::tuple<std::vector<ActsExamples::ConstTrackContainer*>, std::vector<ActsExamples::Trajectories*>>
AmbiguitySolver::process(std::vector<const ActsExamples::ConstTrackContainer*> input_container,
const edm4eic::Measurement2DCollection& meas2Ds) {

// Assuming ActsExamples::ConstTrackContainer is compatible with Acts::ConstVectorTrackContainer
// Create track container
std::vector<ActsExamples::Trajectories*> output_trajectories;
std::vector<ActsExamples::ConstTrackContainer*> output_tracks;

auto& input_trks = input_container.front();
Acts::GreedyAmbiguityResolution::State state;
m_core->computeInitialState(*input_trks, state, &sourceLinkHash, &sourceLinkEquality);
m_core->resolve(state);

ActsExamples::TrackContainer solvedTracks{std::make_shared<Acts::VectorTrackContainer>(),
std::make_shared<Acts::VectorMultiTrajectory>()};
solvedTracks.ensureDynamicColumns(*input_trks);

for (auto iTrack : state.selectedTracks) {

auto destProxy = solvedTracks.getTrack(solvedTracks.addTrack());
auto srcProxy = input_trks->getTrack(state.trackTips.at(iTrack));
destProxy.copyFrom(srcProxy, false);
destProxy.tipIndex() = srcProxy.tipIndex();

}

output_tracks.push_back(new ActsExamples::ConstTrackContainer(
std::make_shared<Acts::ConstVectorTrackContainer>(std::move(solvedTracks.container())),
input_trks->trackStateContainerHolder()));

//Make output trajectories
ActsExamples::Trajectories::IndexedParameters parameters;
std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;

for (const auto& track : *(output_tracks.front())) {

tips.clear();
parameters.clear();

tips.push_back(track.tipIndex());
parameters.emplace(
std::pair{track.tipIndex(),
ActsExamples::TrackParameters{track.referenceSurface().getSharedPtr(),
track.parameters(), track.covariance(),
track.particleHypothesis()}});

output_trajectories.push_back(new ActsExamples::Trajectories(
((*output_tracks.front())).trackStateContainer(),
tips, parameters));

}

return std::make_tuple(std::move(output_tracks), std::move(output_trajectories));
}

} // namespace eicrecon
42 changes: 42 additions & 0 deletions src/algorithms/tracking/AmbiguitySolver.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2024 Minjung Kim, Barak Schmookler
#pragma once

#include <Acts/Utilities/Logger.hpp>
#include <ActsExamples/EventData/Track.hpp>
#include <ActsExamples/EventData/Trajectories.hpp>
#include <edm4eic/Measurement2D.h>
#include <spdlog/logger.h>
#include <memory>
#include <tuple>
#include <vector>

#include "Acts/AmbiguityResolution/GreedyAmbiguityResolution.hpp"
#include "AmbiguitySolverConfig.h"
#include "algorithms/interfaces/WithPodConfig.h"

namespace eicrecon {

/*Reco Track Filtering Based on Greedy ambiguity resolution solver adopted from ACTS*/
class AmbiguitySolver : public WithPodConfig<eicrecon::AmbiguitySolverConfig> {
public:
AmbiguitySolver();

void init(std::shared_ptr<spdlog::logger> log);

std::tuple<
std::vector<ActsExamples::ConstTrackContainer *>,
std::vector<ActsExamples::Trajectories *>
>
process(std::vector<const ActsExamples::ConstTrackContainer*> input_container,const edm4eic::Measurement2DCollection& meas2Ds);

private:
std::shared_ptr<spdlog::logger> m_log;
Acts::GreedyAmbiguityResolution::Config m_acts_cfg;
std::unique_ptr<Acts::GreedyAmbiguityResolution> m_core;
/// Private access to the logging instance
std::shared_ptr<const Acts::Logger> m_acts_logger{nullptr};
const Acts::Logger& logger() const { return *m_acts_logger; }
};

} // namespace eicrecon
17 changes: 17 additions & 0 deletions src/algorithms/tracking/AmbiguitySolverConfig.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2024 Minjung Kim

#pragma once

#include <vector>

namespace eicrecon {
struct AmbiguitySolverConfig {
/// Maximum amount of shared hits per track.
std::uint32_t maximum_shared_hits = 1;
/// Maximum number of iterations
std::uint32_t maximum_iterations = 100000;
/// Minimum number of measurement to form a track.
std::size_t n_measurements_min = 3;
};
} // namespace eicrecon
52 changes: 52 additions & 0 deletions src/global/tracking/AmbiguitySolver_factory.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2024 Minjung Kim, Barak Schmookler
#pragma once

#include "algorithms/tracking/AmbiguitySolver.h"
#include "algorithms/tracking/AmbiguitySolverConfig.h"
#include "extensions/jana/JOmniFactory.h"
#include "extensions/spdlog/SpdlogMixin.h"
#include <ActsExamples/EventData/Track.hpp>
#include <JANA/JEvent.h>
#include <memory>
#include <string>
#include <utility>
#include <vector>

namespace eicrecon {

class AmbiguitySolver_factory
: public JOmniFactory<AmbiguitySolver_factory, AmbiguitySolverConfig> {

private:
using AlgoT = eicrecon::AmbiguitySolver;
std::unique_ptr<AlgoT> m_algo;

Input<ActsExamples::ConstTrackContainer> m_acts_tracks_input {this};
PodioInput<edm4eic::Measurement2D> m_measurements_input {this};
Output<ActsExamples::ConstTrackContainer> m_acts_tracks_output {this};
Output<ActsExamples::Trajectories> m_acts_trajectories_output {this};

ParameterRef<std::uint32_t> m_maximumSharedHits{this, "maximumSharedHits", config().maximum_shared_hits,
"Maximum number of shared hits allowed"};
ParameterRef<std::uint32_t> m_maximumIterations{this, "maximumIterations", config().maximum_iterations,
"Maximum number of iterations"};
ParameterRef<std::size_t> m_nMeasurementsMin{
this, "nMeasurementsMin", config().n_measurements_min,
"Number of measurements required for further reconstruction"};

public:
void Configure() {
m_algo = std::make_unique<AlgoT>();
m_algo->applyConfig(config());
m_algo->init(logger());
}

void ChangeRun(int64_t run_number) {}

void Process(int64_t run_number, uint64_t event_number) {
std::tie(m_acts_tracks_output(),m_acts_trajectories_output()) = m_algo->process(m_acts_tracks_input(),*m_measurements_input());
}
} ;

} // namespace eicrecon
71 changes: 63 additions & 8 deletions src/global/tracking/tracking.cc
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

#include "ActsToTracks.h"
#include "ActsToTracks_factory.h"
#include "AmbiguitySolver_factory.h"
#include "CKFTracking_factory.h"
#include "IterativeVertexFinder_factory.h"
#include "TrackParamTruthInit_factory.h"
Expand Down Expand Up @@ -62,10 +63,10 @@ void InitPlugin(JApplication *app) {
std::vector<std::string> input_collections;
auto readouts = app->GetService<DD4hep_service>()->detector()->readouts();
for (const auto& [hit_collection, rec_collection] : possible_collections) {
if (readouts.find(hit_collection) != readouts.end()) {
// Add the collection to the list of input collections
input_collections.push_back(rec_collection);
}
if (readouts.find(hit_collection) != readouts.end()) {
// Add the collection to the list of input collections
input_collections.push_back(rec_collection);
}
}

// Tracker hits collector
Expand All @@ -89,8 +90,35 @@ void InitPlugin(JApplication *app) {
"CentralTrackerMeasurements"
},
{
"CentralCKFActsTrajectories",
"CentralCKFActsTracks",
"CentralCKFActsTrajectoriesUnfiltered",
"CentralCKFActsTracksUnfiltered",
},
app
));

app->Add(new JOmniFactoryGeneratorT<ActsToTracks_factory>(
"CentralCKFTracksUnfiltered",
{
"CentralTrackerMeasurements",
"CentralCKFActsTrajectoriesUnfiltered",
},
{
"CentralCKFTrajectoriesUnfiltered",
"CentralCKFTrackParametersUnfiltered",
"CentralCKFTracksUnfiltered",
},
app
));

app->Add(new JOmniFactoryGeneratorT<AmbiguitySolver_factory>(
"AmbiguityResolutionSolver",
{
"CentralCKFActsTracksUnfiltered",
"CentralTrackerMeasurements"
},
{
"CentralCKFActsTracks",
"CentralCKFActsTrajectories",
},
app
));
Expand Down Expand Up @@ -124,8 +152,35 @@ void InitPlugin(JApplication *app) {
"CentralTrackerMeasurements"
},
{
"CentralCKFSeededActsTrajectories",
"CentralCKFSeededActsTracks",
"CentralCKFSeededActsTrajectoriesUnfiltered",
"CentralCKFSeededActsTracksUnfiltered",
},
app
));

app->Add(new JOmniFactoryGeneratorT<ActsToTracks_factory>(
"CentralCKFSeededTracksUnfiltered",
{
"CentralTrackerMeasurements",
"CentralCKFSeededActsTrajectoriesUnfiltered",
},
{
"CentralCKFSeededTrajectoriesUnfiltered",
"CentralCKFSeededTrackParametersUnfiltered",
"CentralCKFSeededTracksUnfiltered",
},
app
));

app->Add(new JOmniFactoryGeneratorT<AmbiguitySolver_factory>(
"SeededAmbiguityResolutionSolver",
{
"CentralCKFSeededActsTracksUnfiltered",
"CentralTrackerMeasurements"
},
{
"CentralCKFSeededActsTracks",
"CentralCKFSeededActsTrajectories",
},
app
));
Expand Down
8 changes: 8 additions & 0 deletions src/services/io/podio/JEventProcessorPODIO.cc
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,14 @@ JEventProcessorPODIO::JEventProcessorPODIO() {
"CentralCKFSeededTrajectories",
"CentralCKFSeededTracks",
"CentralCKFSeededTrackParameters",
//tracking properties - true seeding
"CentralCKFTrajectoriesUnfiltered",
"CentralCKFTracksUnfiltered",
"CentralCKFTrackParametersUnfiltered",
//tracking properties - realistic seeding
"CentralCKFSeededTrajectoriesUnfiltered",
"CentralCKFSeededTracksUnfiltered",
"CentralCKFSeededTrackParametersUnfiltered",
"InclusiveKinematicsDA",
"InclusiveKinematicsJB",
"InclusiveKinematicsSigma",
Expand Down

0 comments on commit 1413111

Please sign in to comment.