diff --git a/PWGPP/CMakeLists.txt b/PWGPP/CMakeLists.txt index 7565dace413..9d96b0b8d33 100644 --- a/PWGPP/CMakeLists.txt +++ b/PWGPP/CMakeLists.txt @@ -12,5 +12,5 @@ # add_subdirectory(Core) # add_subdirectory(DataModel) add_subdirectory(Tasks) -# add_subdirectory(TableProducer) +add_subdirectory(TableProducer) diff --git a/PWGPP/TableProducer/CMakeLists.txt b/PWGPP/TableProducer/CMakeLists.txt new file mode 100644 index 00000000000..a747c9f892b --- /dev/null +++ b/PWGPP/TableProducer/CMakeLists.txt @@ -0,0 +1,16 @@ +# Copyright 2019-2020 CERN and copyright holders of ALICE O2. +# See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +# All rights not expressly granted are reserved. +# +# This software is distributed under the terms of the GNU General Public +# License v3 (GPL Version 3), copied verbatim in the file "COPYING". +# +# In applying this license CERN does not waive the privileges and immunities +# granted to it by virtue of its status as an Intergovernmental Organization +# or submit itself to any jurisdiction. + +o2physics_add_dpl_workflow(pid-ml-producer + SOURCES pidMLProducer.cxx + JOB_POOL analysis + PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore + COMPONENT_NAME Analysis) diff --git a/PWGPP/TableProducer/pidMLProducer.cxx b/PWGPP/TableProducer/pidMLProducer.cxx new file mode 100644 index 00000000000..ca4aa06ff57 --- /dev/null +++ b/PWGPP/TableProducer/pidMLProducer.cxx @@ -0,0 +1,128 @@ +// Copyright CERN and copyright holders of ALICE O2. This software is +// distributed under the terms of the GNU General Public License v3 (GPL +// Version 3), copied verbatim in the file "COPYING". +// +// See http://alice-o2.web.cern.ch/license for full licensing information. +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. +#include "Framework/AnalysisTask.h" +#include "Framework/AnalysisDataModel.h" +#include "Common/Core/MC.h" +#include "Common/Core/PID/PIDResponse.h" +#include "Common/DataModel/TrackSelectionTables.h" + +using namespace o2; +using namespace o2::framework; +using namespace o2::framework::expressions; + +namespace o2::aod +{ +namespace pidtracks +{ +DECLARE_SOA_COLUMN(Px, px, float); //! Non-dynamic column with track x-momentum +DECLARE_SOA_COLUMN(Py, py, float); //! Non-dynamic column with track y-momentum +DECLARE_SOA_COLUMN(Pz, pz, float); //! Non-dynamic column with track z-momentum +DECLARE_SOA_COLUMN(Sign, sign, float); //! Non-dynamic column with track sign +DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, uint8_t); //! +} // namespace pidtracks +DECLARE_SOA_TABLE(PidTracksReal, "AOD", "PIDTRACKSREAL", //! Real tracks for prediction and domain adaptation + aod::track::TPCSignal, + aod::pidtofsignal::TOFSignal, + pidtracks::Px, + pidtracks::Py, + pidtracks::Pz, + pidtracks::Sign, + aod::track::X, + aod::track::Y, + aod::track::Z, + aod::track::Alpha, + aod::track::TrackType, + aod::track::TPCNClsShared, + aod::track::DcaXY, + aod::track::DcaZ); +DECLARE_SOA_TABLE(PidTracksMc, "AOD", "PIDTRACKSMC", //! MC tracks for training + aod::track::TPCSignal, + aod::pidtofsignal::TOFSignal, + pidtracks::Px, + pidtracks::Py, + pidtracks::Pz, + pidtracks::Sign, + aod::track::X, + aod::track::Y, + aod::track::Z, + aod::track::Alpha, + aod::track::TrackType, + aod::track::TPCNClsShared, + aod::track::DcaXY, + aod::track::DcaZ, + aod::mcparticle::PdgCode, + pidtracks::IsPhysicalPrimary); +} //namespace o2::aod + +void customize(std::vector& workflowOptions) +{ + ConfigParamSpec optionDoMC{"doMC", VariantType::Bool, false, {"Fill PID train table with MC data."}}; + workflowOptions.push_back(optionDoMC); +} + +#include "Framework/runDataProcessing.h" + +struct CreateTableMc { + Produces pidTracksTable; + + Filter trackFilter = aod::track::isGlobalTrack == (uint8_t) true; + using BigTracksMC = soa::Filtered>; + + void process(BigTracksMC const& tracks, aod::McParticles const& mctracks) + { + for (const auto& track : tracks) { + const auto mcParticle = track.mcParticle(); + uint8_t isPrimary = (uint8_t)MC::isPhysicalPrimary(mcParticle); + pidTracksTable(track.tpcSignal(), track.tofSignal(), + track.px(), track.py(), track.pz(), + track.sign(), + track.x(), track.y(), track.z(), + track.alpha(), + track.trackType(), + track.tpcNClsShared(), + track.dcaXY(), track.dcaZ(), + mcParticle.pdgCode(), + isPrimary); + } + } +}; + +struct CreateTableReal { + Produces pidTracksTable; + + Filter trackFilter = aod::track::isGlobalTrack == (uint8_t) true; + using BigTracks = soa::Filtered>; + + void process(BigTracks const& tracks) + { + for (const auto& track : tracks) { + pidTracksTable(track.tpcSignal(), track.tofSignal(), + track.px(), track.py(), track.pz(), + track.sign(), + track.x(), track.y(), track.z(), + track.alpha(), + track.trackType(), + track.tpcNClsShared(), + track.dcaXY(), track.dcaZ()); + } + } +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + const bool doMC = cfgc.options().get("doMC"); + if (doMC) { + return WorkflowSpec{ + adaptAnalysisTask(cfgc)}; + } else { + return WorkflowSpec{ + adaptAnalysisTask(cfgc)}; + } +}