add basic unit support to RMGAnalysisReader

include/RMGAnalysisReader.hh

@@ -16,6 +16,7 @@
 #ifndef _RMG_ANALYSIS_READER_HH_
 #define _RMG_ANALYSIS_READER_HH_
 
+#include <map>
 #include <string>
 
 #include "G4AutoLock.hh"
@@ -60,6 +61,7 @@ class RMGAnalysisReader final {
         void unlock() {
           fReader = nullptr;
           fNtupleId = -1;
+          fUnits = nullptr;
           if (fLock) { fLock.unlock(); }
         }
 
@@ -68,33 +70,48 @@ 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<std::string>& 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<std::string>& 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<std::string>& 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; }
 
       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<std::string, std::string>* u)
+            : fLock(std::move(lock)), fReader(reader), fNtupleId(nt), fUnits(u) {};
         Access(Access&&) = default;
 
+        void AssertUnit(const std::string& name, const std::vector<std::string>& allowed_units) const;
+
         G4VAnalysisReader* fReader = nullptr;
         int fNtupleId = -1;
+        const std::map<std::string, std::string>* fUnits;
         G4AutoLock fLock;
     };
 
@@ -146,6 +163,9 @@ class RMGAnalysisReader final {
     G4VAnalysisReader* fReader = nullptr;
     int fNtupleId = -1;
 
+    std::map<std::string, std::string> fUnits{};
+    bool fHasUnits = false;
+
     std::string fFileName;
     bool fFileIsTemp = false;
 };

include/RMGConvertLH5.hh

@@ -18,6 +18,7 @@
 #define _RMG_CONVERT_LH5_HH
 
 #include <filesystem>
+#include <map>
 #include <memory>
 #include <optional>
 #include <regex>
@@ -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<std::string, std::map<std::string, std::string>>&);
 
     inline static bool fIsStandalone = false;
 
@@ -77,12 +79,12 @@ class RMGConvertLH5 {
     ////////////////////////////////////////////////////////////////////////////////////////////
     // LH5 -> HDF5 (input files):
 
-    bool ConvertFromLH5Internal();
+    bool ConvertFromLH5Internal(std::map<std::string, std::map<std::string, std::string>>&);
 
     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<std::string, std::string>&);
 
     ////////////////////////////////////////////////////////////////////////////////////////////
 

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 <https://www.gnu.org/licenses/>.
 
+#include "RMGAnalysisReader.hh"
+
 #include <filesystem>
 #include <random>
-
-#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<std::string, std::map<std::string, std::string>> 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<std::string>& 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);
+  }
 }

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<std::string, std::string>& 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<std::string, std::map<std::string, std::string>>& 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<std::string, std::map<std::string, std::string>>& 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;

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<std::string, double> 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);
 }

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<std::string, double> 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;
     }
 

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<std::string, std::map<std::string, std::string>> units_map{};
+    RMGConvertLH5::ConvertFromLH5(file_name, ntuple_group_name, dry_run, !file_names.empty(),
+        units_map);
   }
 
   struct rusage usage{};

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)