From d62242da1a8b42f69fcd108281e8b36361a9ce64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=B2=20Jacazio?= <nicolo.jacazio@cern.ch>
Date: Wed, 13 Apr 2022 12:31:46 +0200
Subject: [PATCH] Add energy as a dynamic column

---
 .../Core/include/Framework/AnalysisDataModel.h   | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/Framework/Core/include/Framework/AnalysisDataModel.h b/Framework/Core/include/Framework/AnalysisDataModel.h
index 72deea106560f..d6cda2dd85009 100644
--- a/Framework/Core/include/Framework/AnalysisDataModel.h
+++ b/Framework/Core/include/Framework/AnalysisDataModel.h
@@ -135,12 +135,19 @@ DECLARE_SOA_DYNAMIC_COLUMN(Pz, pz, //! Momentum in z-direction in GeV/c
 
 DECLARE_SOA_EXPRESSION_COLUMN(P, p, float, //! Momentum in Gev/c
                               0.5f * (ntan(PIQuarter - 0.5f * natan(aod::track::tgl)) + 1.f / ntan(PIQuarter - 0.5f * natan(aod::track::tgl))) / nabs(aod::track::signed1Pt));
+DECLARE_SOA_DYNAMIC_COLUMN(Energy, energy, //! Track energy, computed under the mass assumption given as input
+                           [](float signed1Pt, float tgl, float mass) -> float {
+                             const auto pt = 1.f / std::abs(signed1Pt);
+                             const auto pz = pt * tgl;
+                             const auto p = 0.5f * (tan(PIQuarter - 0.5f * atan(tgl)) + 1.f / tan(PIQuarter - 0.5f * atan(tgl))) * pt;
+                             return sqrt(p * p + mass * mass);
+                           });
 DECLARE_SOA_DYNAMIC_COLUMN(Rapidity, rapidity, //! Track rapidity, computed under the mass assumption given as input
                            [](float signed1Pt, float tgl, float mass) -> float {
-                             auto pt = 1.f / std::abs(signed1Pt);
-                             auto pz = pt * tgl;
-                             auto p = 0.5f * (tan(PIQuarter - 0.5f * atan(tgl)) + 1.f / tan(PIQuarter - 0.5f * atan(tgl))) * pt;
-                             auto energy = sqrt(p * p + mass * mass);
+                             const auto pt = 1.f / std::abs(signed1Pt);
+                             const auto pz = pt * tgl;
+                             const auto p = 0.5f * (tan(PIQuarter - 0.5f * atan(tgl)) + 1.f / tan(PIQuarter - 0.5f * atan(tgl))) * pt;
+                             const auto energy = sqrt(p * p + mass * mass);
                              return 0.5f * log((energy + pz) / (energy - pz));
                            });
 
@@ -282,6 +289,7 @@ DECLARE_SOA_TABLE_FULL(StoredTracks, "Tracks", "AOD", "TRACK", //! On disk versi
                        track::Px<track::Signed1Pt, track::Snp, track::Alpha>,
                        track::Py<track::Signed1Pt, track::Snp, track::Alpha>,
                        track::Pz<track::Signed1Pt, track::Tgl>,
+                       track::Energy<track::Signed1Pt, track::Tgl>,
                        track::Rapidity<track::Signed1Pt, track::Tgl>,
                        track::Sign<track::Signed1Pt>);