From 38ced948b701ca4e1699e15f46eadb2eea4f6dd6 Mon Sep 17 00:00:00 2001 From: Manuel Huber Date: Fri, 24 Jan 2025 15:45:53 +0100 Subject: [PATCH] add basic unit support to RMGAnalysisReader --- include/RMGAnalysisReader.hh | 30 ++++++++++++++--- include/RMGConvertLH5.hh | 8 +++-- src/RMGAnalysisReader.cc | 54 ++++++++++++++++++++++++------ src/RMGConvertLH5.cc | 18 ++++++---- src/RMGGeneratorFromFile.cc | 16 +++++---- src/RMGVertexFromFile.cc | 19 ++++++++--- src/remage-from-lh5.cc | 4 ++- tests/vertex/CMakeLists.txt | 2 +- tests/vertex/create-test-input.py | 30 ++++++++++------- tests/vertex/macros/pos-lh5-mm.mac | 13 +++++++ tests/vertex/run-vtx.py | 32 +++++++++++++++--- 11 files changed, 172 insertions(+), 54 deletions(-) create mode 100644 tests/vertex/macros/pos-lh5-mm.mac diff --git a/include/RMGAnalysisReader.hh b/include/RMGAnalysisReader.hh index 964e8cb..820026d 100644 --- a/include/RMGAnalysisReader.hh +++ b/include/RMGAnalysisReader.hh @@ -16,6 +16,7 @@ #ifndef _RMG_ANALYSIS_READER_HH_ #define _RMG_ANALYSIS_READER_HH_ +#include #include #include "G4AutoLock.hh" @@ -60,6 +61,7 @@ class RMGAnalysisReader final { void unlock() { fReader = nullptr; fNtupleId = -1; + fUnits = nullptr; if (fLock) { fLock.unlock(); } } @@ -68,20 +70,31 @@ class RMGAnalysisReader final { [[nodiscard]] auto GetNtupleRow() { return fReader->GetNtupleRow(fNtupleId); } /** * @brief wraps @ref G4VAnalysisReader::SetNtupleDColumn. */ - auto SetNtupleDColumn(const std::string& name, G4double& value) { + auto SetNtupleDColumn(const std::string& name, G4double& value, + const std::vector& allowed_units = {}) { + AssertUnit(name, allowed_units); return fReader->SetNtupleDColumn(fNtupleId, name, value); } /** * @brief wraps @ref G4VAnalysisReader::SetNtupleFColumn. */ - auto SetNtupleFColumn(const std::string& name, G4float& value) { + auto SetNtupleFColumn(const std::string& name, G4float& value, + const std::vector& allowed_units = {}) { + AssertUnit(name, allowed_units); return fReader->SetNtupleFColumn(fNtupleId, name, value); } /** * @brief wraps @ref G4VAnalysisReader::SetNtupleIColumn. */ - auto SetNtupleIColumn(const std::string& name, G4int& value) { + auto SetNtupleIColumn(const std::string& name, G4int& value, + const std::vector& allowed_units = {}) { + AssertUnit(name, allowed_units); return fReader->SetNtupleIColumn(fNtupleId, name, value); } + /** + * @brief get unit information for the column. an empty string means either no unit + * attached or no support by the file format. */ + [[nodiscard]] std::string GetUnit(const std::string& name) const; + /** * @brief check whether this access handle is still valid. */ operator bool() const { return fReader != nullptr && fNtupleId >= 0 && fLock; } @@ -89,12 +102,16 @@ class RMGAnalysisReader final { private: // only allow creation or moving in parent. - inline Access(G4AutoLock lock, G4VAnalysisReader* reader, int nt) - : fLock(std::move(lock)), fReader(reader), fNtupleId(nt) {}; + inline Access(G4AutoLock lock, G4VAnalysisReader* reader, int nt, + const std::map* u) + : fLock(std::move(lock)), fReader(reader), fNtupleId(nt), fUnits(u) {}; Access(Access&&) = default; + void AssertUnit(const std::string& name, const std::vector& allowed_units) const; + G4VAnalysisReader* fReader = nullptr; int fNtupleId = -1; + const std::map* fUnits; G4AutoLock fLock; }; @@ -146,6 +163,9 @@ class RMGAnalysisReader final { G4VAnalysisReader* fReader = nullptr; int fNtupleId = -1; + std::map fUnits{}; + bool fHasUnits = false; + std::string fFileName; bool fFileIsTemp = false; }; diff --git a/include/RMGConvertLH5.hh b/include/RMGConvertLH5.hh index 0516fa5..cda94f8 100644 --- a/include/RMGConvertLH5.hh +++ b/include/RMGConvertLH5.hh @@ -18,6 +18,7 @@ #define _RMG_CONVERT_LH5_HH #include +#include #include #include #include @@ -33,7 +34,8 @@ class RMGConvertLH5 { public: static bool ConvertToLH5(std::string, std::string, bool, bool part_of_batch = false); - static bool ConvertFromLH5(std::string, std::string, bool, bool part_of_batch = false); + static bool ConvertFromLH5(std::string, std::string, bool, bool part_of_batch, + std::map>&); inline static bool fIsStandalone = false; @@ -77,12 +79,12 @@ class RMGConvertLH5 { //////////////////////////////////////////////////////////////////////////////////////////// // LH5 -> HDF5 (input files): - bool ConvertFromLH5Internal(); + bool ConvertFromLH5Internal(std::map>&); static inline const std::regex table_dtype_re = std::regex("^table\\{.*\\}$"); std::string HDFDataTypeToForm(H5::DataType); - bool ConvertTableToNTuple(H5::Group&); + bool ConvertTableToNTuple(H5::Group&, std::map&); //////////////////////////////////////////////////////////////////////////////////////////// diff --git a/src/RMGAnalysisReader.cc b/src/RMGAnalysisReader.cc index 369429e..e795a49 100644 --- a/src/RMGAnalysisReader.cc +++ b/src/RMGAnalysisReader.cc @@ -13,26 +13,24 @@ // You should have received a copy of the GNU Lesser General Public License // along with this program. If not, see . +#include "RMGAnalysisReader.hh" + #include #include - -#include "RMGVertexFromFile.hh" namespace fs = std::filesystem; #include "G4AnalysisUtilities.hh" -#include "G4RootAnalysisReader.hh" -#include "G4Threading.hh" -#include "G4VAnalysisReader.hh" +#include "G4CsvAnalysisReader.hh" #if RMG_HAS_HDF5 #include "G4Hdf5AnalysisReader.hh" #endif -#include "G4CsvAnalysisReader.hh" +#include "G4RootAnalysisReader.hh" +#include "G4Threading.hh" #include "G4XmlAnalysisReader.hh" #if RMG_HAS_HDF5 #include "RMGConvertLH5.hh" #endif - #include "RMGLog.hh" G4Mutex RMGAnalysisReader::fMutex = G4MUTEX_INITIALIZER; @@ -43,7 +41,7 @@ RMGAnalysisReader::Access RMGAnalysisReader::OpenFile(std::string& file_name, // create an invalid lock to return on error. G4AutoLock invalid_lock(&fMutex, std::defer_lock); invalid_lock.release(); - auto invalid_access = RMGAnalysisReader::Access(std::move(invalid_lock), nullptr, -1); + auto invalid_access = RMGAnalysisReader::Access(std::move(invalid_lock), nullptr, -1, nullptr); // reader initialization should only happen on the master thread (otherwise it will fail). if (!G4Threading::IsMasterThread()) { @@ -62,6 +60,8 @@ RMGAnalysisReader::Access RMGAnalysisReader::OpenFile(std::string& file_name, return std::move(invalid_access); } + std::map> units_map{}; + // NOTE: GetExtension() returns a default extension if there is no file extension auto ext = G4Analysis::GetExtension(file_name); if (ext == "root") { @@ -88,12 +88,13 @@ RMGAnalysisReader::Access RMGAnalysisReader::OpenFile(std::string& file_name, return std::move(invalid_access); } - if (!RMGConvertLH5::ConvertFromLH5(new_fn, ntuple_dir_name, false)) { + if (!RMGConvertLH5::ConvertFromLH5(new_fn, ntuple_dir_name, false, false, units_map)) { RMGLog::Out(RMGLog::error, "Conversion of input file ", new_fn, " to LH5 failed. Data is potentially corrupted."); return std::move(invalid_access); } + fHasUnits = true; file_name = new_fn; fReader = G4Hdf5AnalysisReader::Instance(); fFileIsTemp = true; @@ -118,7 +119,8 @@ RMGAnalysisReader::Access RMGAnalysisReader::OpenFile(std::string& file_name, RMGLog::Out(RMGLog::error, "Ntuple named '", ntuple_name, "' could not be found in input file!"); return std::move(invalid_access); } - return {std::move(lock), fReader, fNtupleId}; + fUnits = units_map[ntuple_name]; + return {std::move(lock), fReader, fNtupleId, fHasUnits ? &fUnits : nullptr}; } void RMGAnalysisReader::CloseFile() { @@ -141,11 +143,41 @@ void RMGAnalysisReader::CloseFile() { fFileName = ""; fFileIsTemp = false; fNtupleId = -1; + fHasUnits = false; + fHasUnits = {}; } RMGAnalysisReader::Access RMGAnalysisReader::GetLockedReader() const { G4AutoLock lock(&fMutex); - return {std::move(lock), fReader, fNtupleId}; + return {std::move(lock), fReader, fNtupleId, fHasUnits ? &fUnits : nullptr}; +} + +G4AutoLock RMGAnalysisReader::GetLock() const { + + // reader initialization should only happen on the master thread (otherwise it will fail). + if (!G4Threading::IsMasterThread()) { + RMGLog::OutDev(RMGLog::fatal, "can only be used on the master thread"); + std::abort(); + } + + G4AutoLock lock(&fMutex); + + return std::move(lock); +} + +/* ========================================================================================== */ + +std::string RMGAnalysisReader::Access::GetUnit(const std::string& name) const { + if (!fUnits) return ""; + return fUnits->at(name); +} + +void RMGAnalysisReader::Access::AssertUnit(const std::string& name, + const std::vector& allowed_units) const { + if (!fUnits) return; + if (std::find(allowed_units.begin(), allowed_units.end(), GetUnit(name)) == allowed_units.end()) { + RMGLog::Out(RMGLog::fatal, "invalid unit '", GetUnit(name), "' for column ", name); + } } diff --git a/src/RMGConvertLH5.cc b/src/RMGConvertLH5.cc index 9fcee71..2dd8ebd 100644 --- a/src/RMGConvertLH5.cc +++ b/src/RMGConvertLH5.cc @@ -345,7 +345,8 @@ std::string RMGConvertLH5::HDFDataTypeToForm(H5::DataType dtype) { } } -bool RMGConvertLH5::ConvertTableToNTuple(H5::Group& det_group) { +bool RMGConvertLH5::ConvertTableToNTuple(H5::Group& det_group, + std::map& units_map) { const std::string ntuple_name = det_group.getObjName(); const std::string ntuple_log_prefix = "ntuple " + ntuple_name + " - "; LH5Log(RMGLog::detail, ntuple_log_prefix, "visiting"); @@ -381,7 +382,9 @@ bool RMGConvertLH5::ConvertTableToNTuple(H5::Group& det_group) { auto dset_column = det_group.openDataSet(lgdo_name); if (dset_column.attrExists("units")) { - column += "_in_" + GetStringAttribute(dset_column, "units").value(); + units_map[lgdo_name] = GetStringAttribute(dset_column, "units").value(); + } else { + units_map[lgdo_name] = ""; } names_string += column + std::string("\0", 1); forms_string += HDFDataTypeToForm(dset_column.getDataType()) + std::string("\0", 1); @@ -426,7 +429,8 @@ bool RMGConvertLH5::ConvertTableToNTuple(H5::Group& det_group) { return true; } -bool RMGConvertLH5::ConvertFromLH5Internal() { +bool RMGConvertLH5::ConvertFromLH5Internal( + std::map>& units_map) { // using the core driver with no backing storage will allow to change the file purely in-memory. // warning: this will internally allocate approx. the full file size! H5::FileAccPropList fapl; @@ -449,7 +453,8 @@ bool RMGConvertLH5::ConvertFromLH5Internal() { bool ntuple_success = true; for (auto& ntuple : ntuples) { auto det_group = ntuples_group.openGroup(ntuple); - ntuple_success &= ConvertTableToNTuple(det_group); + units_map[ntuple] = {}; + ntuple_success &= ConvertTableToNTuple(det_group, units_map[ntuple]); } if (ntuples_group.attrExists("datatype")) ntuples_group.removeAttr("datatype"); @@ -462,10 +467,11 @@ bool RMGConvertLH5::ConvertFromLH5Internal() { } bool RMGConvertLH5::ConvertFromLH5(std::string lh5_file_name, std::string ntuple_group_name, - bool dry_run, bool part_of_batch) { + bool dry_run, bool part_of_batch, + std::map>& units_map) { auto conv = RMGConvertLH5(lh5_file_name, ntuple_group_name, dry_run, part_of_batch); try { - return conv.ConvertFromLH5Internal(); + return conv.ConvertFromLH5Internal(units_map); } catch (const H5::Exception& e) { conv.LH5Log(RMGLog::error, e.getDetailMsg()); return false; diff --git a/src/RMGGeneratorFromFile.cc b/src/RMGGeneratorFromFile.cc index 8f68ba1..d8794d0 100644 --- a/src/RMGGeneratorFromFile.cc +++ b/src/RMGGeneratorFromFile.cc @@ -41,11 +41,11 @@ void RMGGeneratorFromFile::OpenFile(std::string& name) { if (!reader) return; // bind the static variables once here, and only use them later. - reader.SetNtupleIColumn("g4_pid", fRowData.fG4Pid); - reader.SetNtupleDColumn("ekin_in_MeV", fRowData.fEkin); - reader.SetNtupleDColumn("px", fRowData.fPx); - reader.SetNtupleDColumn("py", fRowData.fPy); - reader.SetNtupleDColumn("pz", fRowData.fPz); + reader.SetNtupleIColumn("g4_pid", fRowData.fG4Pid, {""}); + reader.SetNtupleDColumn("ekin", fRowData.fEkin, {"eV", "keV", "MeV", "GeV"}); + reader.SetNtupleDColumn("px", fRowData.fPx, {""}); + reader.SetNtupleDColumn("py", fRowData.fPy, {""}); + reader.SetNtupleDColumn("pz", fRowData.fPz, {""}); } void RMGGeneratorFromFile::BeginOfRunAction(const G4Run*) { @@ -84,6 +84,7 @@ void RMGGeneratorFromFile::GeneratePrimaries(G4Event* event) { // make copy of data and exit critical section. auto row_data = fRowData; + auto unit_name = locked_reader.GetUnit("ekin"); locked_reader.unlock(); // check for NaN sentinel values - i.e. non-existing columns (there is no error message). @@ -104,10 +105,13 @@ void RMGGeneratorFromFile::GeneratePrimaries(G4Event* event) { G4ThreeVector momentum{row_data.fPx, row_data.fPy, row_data.fPz}; + const std::map units = {{"", u::MeV}, {"eV", u::eV}, {"keV", u::keV}, + {"MeV", u::MeV}, {"GeV", u::GeV}}; + fGun->SetParticleDefinition(particle); fGun->SetParticlePosition(fParticlePosition); fGun->SetParticleMomentumDirection(momentum); - fGun->SetParticleEnergy(row_data.fEkin * u::MeV); + fGun->SetParticleEnergy(row_data.fEkin * units.at(unit_name)); fGun->GeneratePrimaryVertex(event); } diff --git a/src/RMGVertexFromFile.cc b/src/RMGVertexFromFile.cc index 4054e91..18c829b 100644 --- a/src/RMGVertexFromFile.cc +++ b/src/RMGVertexFromFile.cc @@ -33,9 +33,16 @@ void RMGVertexFromFile::OpenFile(std::string& name) { if (!reader) return; // bind the static variables once here, and only use them later. - reader.SetNtupleDColumn("xloc_in_m", fXpos); - reader.SetNtupleDColumn("yloc_in_m", fYpos); - reader.SetNtupleDColumn("zloc_in_m", fZpos); + reader.SetNtupleDColumn("xloc", fXpos, {"nm", "um", "mm", "cm", "m"}); + reader.SetNtupleDColumn("yloc", fYpos, {"nm", "um", "mm", "cm", "m"}); + reader.SetNtupleDColumn("zloc", fZpos, {"nm", "um", "mm", "cm", "m"}); + + const auto xunit = reader.GetUnit("xloc"); + const auto yunit = reader.GetUnit("yloc"); + const auto zunit = reader.GetUnit("zloc"); + if (xunit != yunit || xunit != zunit) { + RMGLog::OutFormat(RMGLog::fatal, " ('{}', '{}', '{}')", xunit, yunit, zunit); + } } bool RMGVertexFromFile::GenerateVertex(G4ThreeVector& vertex) { @@ -53,7 +60,11 @@ bool RMGVertexFromFile::GenerateVertex(G4ThreeVector& vertex) { return false; } - vertex = G4ThreeVector{fXpos, fYpos, fZpos} * CLHEP::m; + const auto unit_name = reader.GetUnit("xloc"); + const std::map units = {{"", CLHEP::m}, {"nm", CLHEP::nm}, + {"um", CLHEP::um}, {"mm", CLHEP::mm}, {"cm", CLHEP::cm}, {"m", CLHEP::m}}; + + vertex = G4ThreeVector{fXpos, fYpos, fZpos} * units.at(unit_name); return true; } diff --git a/src/remage-from-lh5.cc b/src/remage-from-lh5.cc index d2de39e..9ef9470 100644 --- a/src/remage-from-lh5.cc +++ b/src/remage-from-lh5.cc @@ -50,7 +50,9 @@ int main(int argc, char** argv) { RMGLog::OutFormat(RMGLog::error, "{} does not exist", file_name); continue; } - RMGConvertLH5::ConvertFromLH5(file_name, ntuple_group_name, dry_run, !file_names.empty()); + std::map> units_map{}; + RMGConvertLH5::ConvertFromLH5(file_name, ntuple_group_name, dry_run, !file_names.empty(), + units_map); } struct rusage usage{}; diff --git a/tests/vertex/CMakeLists.txt b/tests/vertex/CMakeLists.txt index 6aefe66..6689ff6 100644 --- a/tests/vertex/CMakeLists.txt +++ b/tests/vertex/CMakeLists.txt @@ -17,7 +17,7 @@ endforeach() add_test(NAME vertex/gen-input COMMAND ${PYTHONPATH} create-test-input.py ${_remage_from_lh5}) set_tests_properties(vertex/gen-input PROPERTIES LABELS extra FIXTURES_SETUP vertex-input-fixture) -set(_macros pos-hdf5 pos-lh5 kin-lh5 kin-hdf5 pos-kin-lh5) +set(_macros pos-hdf5 pos-lh5 pos-lh5-mm kin-lh5 kin-hdf5 pos-kin-lh5) foreach(_mac ${_macros}) add_test(NAME vertex/${_mac} COMMAND ${PYTHONPATH} run-vtx.py ${REMAGE_PYEXE} ${_mac} 1) diff --git a/tests/vertex/create-test-input.py b/tests/vertex/create-test-input.py index eafed34..0b554b3 100755 --- a/tests/vertex/create-test-input.py +++ b/tests/vertex/create-test-input.py @@ -25,25 +25,31 @@ def _convert_lh5_to_hdf5(lh5_fn): assert ret.returncode == 0 -INPUT_FILE_ROWS = int(1e6) +INPUT_FILE_ROWS = int(1e5) rng = np.random.default_rng(123456) # ================================================== # position input -xloc = Array(np.linspace(0, 1000, INPUT_FILE_ROWS), attrs={"units": "m"}) -yloc = Array(rng.uniform(-1, 1, INPUT_FILE_ROWS), attrs={"units": "m"}) -zloc = Array(rng.uniform(-1, 1, INPUT_FILE_ROWS), attrs={"units": "m"}) -vtx_pos_file = "macros/vtx-pos.lh5" -lh5.write( - Table({"xloc": xloc, "yloc": yloc, "zloc": zloc}), - "vtx/pos", - lh5_file=vtx_pos_file, - wo_mode="of", -) -_convert_lh5_to_hdf5(vtx_pos_file) +def _create_pos_input(dtype, vtx_pos_file, units): + attrs = {"units": units} + xloc = Array(np.linspace(0, 1000, INPUT_FILE_ROWS).astype(dtype), attrs=attrs) + yloc = Array(rng.uniform(-1, 1, INPUT_FILE_ROWS).astype(dtype), attrs=attrs) + zloc = Array(rng.uniform(-1, 1, INPUT_FILE_ROWS).astype(dtype), attrs=attrs) + + lh5.write( + Table({"xloc": xloc, "yloc": yloc, "zloc": zloc}), + "vtx/pos", + lh5_file=vtx_pos_file, + wo_mode="of", + ) + _convert_lh5_to_hdf5(vtx_pos_file) + +_create_pos_input(np.float64, "macros/vtx-pos.lh5", "m") +# _create_pos_input(np.float32, "macros/vtx-pos-float32.lh5", "m") +_create_pos_input(np.float64, "macros/vtx-pos-mm.lh5", "mm") # ================================================== # kinetcis input. diff --git a/tests/vertex/macros/pos-lh5-mm.mac b/tests/vertex/macros/pos-lh5-mm.mac new file mode 100644 index 0000000..5e0ef1c --- /dev/null +++ b/tests/vertex/macros/pos-lh5-mm.mac @@ -0,0 +1,13 @@ +/control/execute macros/_init.mac + +/RMG/Generator/Confine FromFile + +/RMG/Generator/Confinement/FromFile/FileName macros/vtx-pos-mm.lh5 + +/RMG/Generator/Select GPS +/gps/particle ion +/gps/ion 81 208 +/gps/energy 0 keV +/gps/ang/type iso + +/run/beamOn 100000 diff --git a/tests/vertex/run-vtx.py b/tests/vertex/run-vtx.py index b8874f9..f4c3010 100755 --- a/tests/vertex/run-vtx.py +++ b/tests/vertex/run-vtx.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re import subprocess import sys from pathlib import Path @@ -20,6 +21,18 @@ for old_f in Path().glob(f"{output_stem}*.lh5"): Path(old_f).unlink() +# "parse" the macro. +pos_input_file = None +kin_input_file = None +with Path(f"macros/{macro}.mac").open("r") as macro_file: + for line in macro_file: + m = re.match(r"/RMG/Generator/FromFile/FileName (.*)\n$", line) + if m: + kin_input_file = m.group(1) + m = re.match(r"/RMG/Generator/Confinement/FromFile/FileName (.*)\n$", line) + if m: + pos_input_file = m.group(1) + # run remage, produce lh5 output. cmd = [ rmg, @@ -47,28 +60,37 @@ # validate output against input files. if "pos" in macro: - input_pos = lh5.read("vtx/pos", lh5_file="macros/vtx-pos.lh5").view_as("pd") + pos_input_file = pos_input_file.replace(".hdf5", ".lh5") + input_pos = lh5.read("vtx/pos", lh5_file=pos_input_file).view_as("pd") output_pos = lh5.read("stp/vertices", lh5_file=files).view_as("pd") output_pos = output_pos.sort_values("xloc") # sort by linear column. output_pos = output_pos[["xloc", "yloc", "zloc"]] uniq, cnt = np.unique(output_pos["xloc"], return_counts=True) if not np.all(cnt <= 1): - raise ValueError(uniq[cnt > 1]) + msg = f"non-unique pos 'indices' {uniq[cnt > 1]}" + raise ValueError(msg) + + input_pos = input_pos.iloc[0 : len(output_pos)] + input_pos = input_pos[["xloc", "yloc", "zloc"]] - input_pos = input_pos.iloc[0 : len(output_pos)][["xloc", "yloc", "zloc"]] + # re-scale to accommodate for different units. + if "vtx-pos-mm.lh5" in pos_input_file: + input_pos /= 1000 assert np.all(np.isclose(output_pos.to_numpy(), input_pos.to_numpy())) if "kin" in macro: - input_kin = lh5.read("vtx/kin", lh5_file="macros/vtx-kin.lh5").view_as("pd") + kin_input_file = kin_input_file.replace(".hdf5", ".lh5") + input_kin = lh5.read("vtx/kin", lh5_file=kin_input_file).view_as("pd") output_kin = lh5.read("stp/particles", lh5_file=files).view_as("pd") output_kin = output_kin.sort_values("ekin") # sort by linear column. output_kin = output_kin[["px", "py", "pz", "ekin", "particle"]] uniq, cnt = np.unique(output_kin["px"], return_counts=True) if not np.all(cnt <= 1): - raise ValueError(uniq[cnt > 1]) + msg = f"non-unique kin 'indices' {uniq[cnt > 1]}" + raise ValueError(msg) output_p_scale = np.sqrt( output_kin["px"] ** 2 + output_kin["py"] ** 2 + output_kin["pz"] ** 2