StandaloneDataIO.h

Go to the documentation of this file.

//
// Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
//
// Dear emacs, this is -*- c++ -*-
//
 
#ifndef CALORECGPU_STANDALONEDATAIO_H
#define CALORECGPU_STANDALONEDATAIO_H
 
#include "Helpers.h"
 
#include "CUDAFriendlyClasses.h"
 
#include <fstream>
#include <string>
#include <map>
#include <set>
#include <iomanip>
#include <sstream>
 
#include <filesystem>
 
struct StandaloneDataIO
{
  enum class ErrorState
  {
    OK = 0, ReadError, WriteError
  };
 
 protected:
 
  inline static void report_error(const std::filesystem::path & file, const std::string & kind, const bool really_report = false)
  {
    if (really_report)
      {
        std::cerr << "ERROR: when " << kind << " from '" << file << "'." << std::endl;
      }
  }
 
 public:
 
  struct ConstantInformation
  {
    friend struct StandaloneDataIO;
 
   public:
 
    inline static ErrorState read_geometry(const std::filesystem::path & file,
                                           CaloRecGPU::Helpers::CPU_object<CaloRecGPU::GeometryArr> & geo,
                                           const bool report = false)
    {
      std::ifstream in(file.native(), std::ios_base::binary);
      geo.binary_input(in);
      if (in.fail())
        {
          report_error(file, "reading geometry", report);
          return ErrorState::ReadError;
        }
      in.close();
      return ErrorState::OK;
    }
 
    inline static ErrorState read_noise(const std::filesystem::path & file,
                                        CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellNoiseArr> & noise,
                                        const bool report = false)
    {
      std::ifstream in(file.native(), std::ios_base::binary);
      noise.binary_input(in);
      if (in.fail())
        {
          report_error(file, "reading noise", report);
          return ErrorState::ReadError;
        }
      in.close();
      return ErrorState::OK;
    }
 
    inline static ErrorState write_geometry(std::filesystem::path file,
                                            const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::GeometryArr> & geo,
                                            const bool report = false)
    {
      file.replace_extension(".geometry");
      std::ofstream out(file, std::ios_base::binary);
      geo.binary_output(out);
      if (out.fail())
        {
          report_error(file, "writing geometry", report);
          return ErrorState::WriteError;
        }
      out.close();
      return ErrorState::OK;
    }
 
    inline static ErrorState write_noise(std::filesystem::path file,
                                         const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellNoiseArr> & noise,
                                         const bool report = false)
    {
      file.replace_extension(".noise");
      std::ofstream out(file, std::ios_base::binary);
      noise.binary_output(out);
      if (out.fail())
        {
          report_error(file, "writing noise", report);
          return ErrorState::WriteError;
        }
      out.close();
      return ErrorState::OK;
    }
  };
 
  struct EventInformation
  {
    friend struct StandaloneDataIO;
 
   public:
    inline static ErrorState read_cluster_info(const std::filesystem::path & file,
                                               CaloRecGPU::Helpers::CPU_object<CaloRecGPU::ClusterInfoArr> & clusters,
                                               const bool report = false)
    {
      clusters.allocate();
 
      std::ifstream in(file.native(), std::ios_base::binary);
 
      auto reader = [&](auto * ptr, const auto & size)
      {
        in.read(reinterpret_cast<char *>(ptr), sizeof(*ptr) * size);
      };
 
      reader(static_cast<CaloRecGPU::ClusterBaseInfo *>(clusters), 1);
 
      if (in.fail() || clusters->number < 0 || clusters->number > CaloRecGPU::NMaxClusters)
        {
          report_error(file, "reading clusters", report);
          return ErrorState::ReadError;
        }
      
      if (clusters->has_basic_info())
        {
          reader(clusters->clusterEnergy, clusters->number);
          reader(clusters->clusterEt,     clusters->number);
          reader(clusters->clusterEta,    clusters->number);
          reader(clusters->clusterPhi,    clusters->number);
          reader(clusters->seedCellIndex, clusters->number);
        }
 
      switch (clusters->state)
        {
          case CaloRecGPU::ClusterInformationState::Tags:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::TagsWithBasicInfo:
            reader(clusters->cells.tags, CaloRecGPU::NCaloCells);
            break;
          case CaloRecGPU::ClusterInformationState::Full:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::WithBasicInfo:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::WithMoments:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::WithExtraMoments:
            reader(clusters->cellsPrefixSum, clusters->number + 1);
            reader(clusters->cells.indices,  clusters->number_cells);
            reader(clusters->cellWeights,    clusters->number_cells);
            reader(clusters->clusterIndices, clusters->number_cells);
            break;
          default:
            break;
        }
 
      if (clusters->has_moments())
        {
          clusters->for_all_moments([&](auto & arr)
          {
            unsigned long long base_size;
            if constexpr (std::is_pointer_v<std::decay_t<decltype(*arr)>>)
              {
                base_size = sizeof(**arr) * sizeof(arr)/sizeof(*arr);
              }
            else
              {
                base_size = sizeof(*arr);
              }
            in.read(reinterpret_cast<char *>(arr), base_size * clusters->number);
          });
        }
 
      if (in.fail())
        {
          report_error(file, "reading clusters", report);
          return ErrorState::ReadError;
        }
 
      in.close();
 
      return ErrorState::OK;
 
    }
 
    //Let's assume we're dealing with a full cell collection,
    //as that is going to be the most common case anyway.
    inline static ErrorState read_cell_info(const std::filesystem::path & file,
                                            CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellInfoArr> & cell_info,
                                            const bool report = false)
    {
      std::ifstream in(file.native(), std::ios_base::binary);
      cell_info.binary_input(in);
      if (in.fail())
        {
          report_error(file, "reading cell info", report);
          return ErrorState::ReadError;
        }
      in.close();
      return ErrorState::OK;
    }
 
    inline static ErrorState write_cluster_info(std::filesystem::path file,
                                                const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::ClusterInfoArr> & clusters,
                                                const bool report = false)
    {
      file.replace_extension(".clusterinfo");
      std::ofstream out(file, std::ios_base::binary);
 
      auto writer = [&](const auto * ptr, const auto & size)
      {
        out.write(reinterpret_cast<const char *>(ptr), sizeof(*ptr) * size);
      };
 
 
      out.write(reinterpret_cast<const char *>(static_cast<const CaloRecGPU::ClusterBaseInfo *>(clusters)), sizeof(CaloRecGPU::ClusterBaseInfo));
 
      if (clusters->has_basic_info())
        {
          writer(clusters->clusterEnergy, clusters->number);
          writer(clusters->clusterEt,     clusters->number);
          writer(clusters->clusterEta,    clusters->number);
          writer(clusters->clusterPhi,    clusters->number);
          writer(clusters->seedCellIndex, clusters->number);
        }
      
      switch (clusters->state)
        {
          case CaloRecGPU::ClusterInformationState::Tags:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::TagsWithBasicInfo:
            writer(clusters->cells.tags, CaloRecGPU::NCaloCells);
            break;
          case CaloRecGPU::ClusterInformationState::Full:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::WithBasicInfo:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::WithMoments:
            [[fallthrough]];
          case CaloRecGPU::ClusterInformationState::WithExtraMoments:
            writer(clusters->cellsPrefixSum, clusters->number + 1);
            writer(clusters->cells.indices,  clusters->number_cells);
            writer(clusters->cellWeights,    clusters->number_cells);
            writer(clusters->clusterIndices, clusters->number_cells);
            break;
          default:
            break;
        }
 
      if (clusters->has_moments())
        {
          clusters->for_all_moments([&](const auto & arr)
          {
            unsigned long long base_size;
            if constexpr (std::is_pointer_v<std::decay_t<decltype(*arr)>>)
              {
                base_size = sizeof(**arr) * sizeof(arr)/sizeof(*arr);
              }
            else
              {
                base_size = sizeof(*arr);
              }
            out.write(reinterpret_cast<const char *>(arr), base_size * clusters->number);
          });
        }
 
      if (out.fail())
        {
          report_error(file, "writing clusters", report);
          return ErrorState::WriteError;
        }
      out.close();
      return ErrorState::OK;
    }
 
    //Let's assume we're dealing with a full cell collection,
    //as that is going to be the most common case anyway.
    inline static ErrorState write_cell_info(std::filesystem::path file,
                                             const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellInfoArr> & cell_info,
                                             const bool report = false)
    {
      file.replace_extension(".cellinfo");
      std::ofstream out(file, std::ios_base::binary);
      cell_info.binary_output(out);
      if (out.fail())
        {
          report_error(file, "writing cell info", report);
          return ErrorState::WriteError;
        }
      out.close();
      return ErrorState::OK;
    }
 
  };
 
 protected:
 
  inline static bool create_or_check_folder(const std::filesystem::path & folder, const bool output_errors = true)
  {
    if (!std::filesystem::exists(folder))
      {
        if (!std::filesystem::create_directory(folder))
          {
            if (output_errors)
              {
                std::cout << "ERROR: folder '" << folder << "' could not be created." << std::endl;
              }
            return false;
          }
      }
    else if (!std::filesystem::is_directory(folder))
      {
        if (output_errors)
          {
            std::cout << "ERROR: folder '" << folder << "' is not a valid folder." << std::endl;
          }
        return false;
      }
    return true;
  }
 
 public:
 
  inline static ErrorState prepare_folder_for_output(const std::filesystem::path & folder, const bool output_errors = true)
  {
    if (!create_or_check_folder(folder, output_errors))
      {
        return ErrorState::WriteError;
      }
    return ErrorState::OK;
  }
 
  inline static std::string build_filename( const std::string & prefix,
                                            const std::string & text,
                                            const std::string & suffix,
                                            const std::string & ext      )
  {
    return prefix + (prefix.size() > 0 ? "_" : "") + text +
           (suffix.size() > 0 ? "_" : "") + suffix + "." + ext;
  }
 
  inline static std::string build_filename( const std::string & prefix,
                                            const size_t event_number,
                                            const std::string & suffix,
                                            const std::string & ext,
                                            const unsigned int num_width = 9 )
  {
    std::ostringstream event_ID_format;
    event_ID_format << std::setfill('0') << std::setw(num_width) << event_number;
    const std::string event_ID = event_ID_format.str();
    return build_filename(prefix, event_ID, suffix, ext);
  }
 
 
  inline static ErrorState save_constants_to_folder(const std::filesystem::path & folder,
                                                    const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::GeometryArr> & geo,
                                                    const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellNoiseArr> & noise,
                                                    const std::string & prefix = "",
                                                    const std::string & suffix = "",
                                                    const bool output_errors = true)
  {
    if (!create_or_check_folder(folder, output_errors))
      {
        return ErrorState::WriteError;
      }
 
    auto filename = [&] (const std::string & text, const std::string & ext)
    {
      return folder / build_filename(prefix, text, suffix, ext);
    };
 
    if (ConstantInformation::write_geometry(filename("geometry", "geo"), geo) != ErrorState::OK)
      {
        return ErrorState::WriteError;
      }
    if (ConstantInformation::write_noise(filename("noise", "noise"), noise) != ErrorState::OK)
      {
        return ErrorState::WriteError;
      }
    return ErrorState::OK;
  }
 
  inline static ErrorState save_event_to_folder(const size_t event_number,
                                                const std::filesystem::path & folder,
                                                const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellInfoArr> & cell_info,
                                                const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::ClusterInfoArr> & clusters,
                                                const std::string & prefix = "",
                                                const std::string & suffix = "",
                                                const unsigned int num_width = 9,
                                                const bool output_errors = true)
  {
    if (!create_or_check_folder(folder, output_errors))
      {
        return ErrorState::WriteError;
      }
 
    std::ostringstream event_ID_format;
    event_ID_format << std::setfill('0') << std::setw(num_width) << event_number;
    const std::string event_ID = event_ID_format.str();
 
    auto filename = [&] (const std::string & ext)
    {
      return folder / build_filename(prefix, event_ID, suffix, ext);
    };
 
    if (EventInformation::write_cell_info(filename("cellinfo"), cell_info) != ErrorState::OK)
      {
        return ErrorState::WriteError;
      }
    if (EventInformation::write_cluster_info(filename("clusterinfo"), clusters) != ErrorState::OK)
      {
        return ErrorState::WriteError;
      }
    return ErrorState::OK;
  }
 
  inline static ErrorState save_cell_info_to_folder(const size_t event_number,
                                                    const std::filesystem::path & folder,
                                                    const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellInfoArr> & cell_info,
                                                    const std::string & prefix = "",
                                                    const std::string & suffix = "",
                                                    const unsigned int num_width = 9,
                                                    const bool output_errors = true)
  {
    if (!create_or_check_folder(folder, output_errors))
      {
        return ErrorState::WriteError;
      }
 
    std::ostringstream event_ID_format;
    event_ID_format << std::setfill('0') << std::setw(num_width) << event_number;
    const std::string event_ID = event_ID_format.str();
 
    auto filename = [&] (const std::string & ext)
    {
      return folder / build_filename(prefix, event_ID, suffix, ext);
    };
 
    if (EventInformation::write_cell_info(filename("cellinfo"), cell_info) != ErrorState::OK)
      {
        return ErrorState::WriteError;
      }
    return ErrorState::OK;
  }
 
  inline static ErrorState save_clusters_to_folder(const size_t event_number,
                                                   const std::filesystem::path & folder,
                                                   const CaloRecGPU::Helpers::CPU_object<CaloRecGPU::ClusterInfoArr> & clusters,
                                                   const std::string & prefix = "",
                                                   const std::string & suffix = "",
                                                   const unsigned int num_width = 9,
                                                   const bool output_errors = true)
  {
    if (!create_or_check_folder(folder, output_errors))
      {
        return ErrorState::WriteError;
      }
 
    std::ostringstream event_ID_format;
    event_ID_format << std::setfill('0') << std::setw(num_width) << event_number;
    const std::string event_ID = event_ID_format.str();
 
    auto filename = [&] (const std::string & ext)
    {
      return folder / build_filename(prefix, event_ID, suffix, ext);
    };
 
    if (EventInformation::write_cluster_info(filename("clusters"), clusters) != ErrorState::OK)
      {
        return ErrorState::WriteError;
      }
    return ErrorState::OK;
  }
 
  struct FolderLoad
  {
    std::map<std::string, CaloRecGPU::Helpers::CPU_object<CaloRecGPU::GeometryArr>> geometry;
    std::map<std::string, CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellNoiseArr>> noise;
    std::map<std::string, CaloRecGPU::Helpers::CPU_object<CaloRecGPU::ClusterInfoArr>> cluster_info;
    std::map<std::string, CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellInfoArr>> cell_info;
  };

  struct FolderLoadOptions
  {
    bool load_cluster_info = false,
         load_cell_info = false,
         load_geometry = false,
         load_noise = false;
 
    static constexpr FolderLoadOptions None()
    {
      FolderLoadOptions ret;
      return ret;
      //By design, all options are initialized as false...
    }
    static constexpr FolderLoadOptions All()
    {
      FolderLoadOptions ret{true, true, true, true};
      return ret;
    }
  };
 

  template <class F>
  inline static FolderLoad load_folder_filter(F && filter_function,
                                              //Receives a std::string (the filename),
                                              //returns `true` if the file should be filtered out.
                                              const std::filesystem::path & folder,
                                              int max_events = -1,
                                              const FolderLoadOptions & flo = FolderLoadOptions::None(),
                                              const bool output_messages = true)
  {
    FolderLoad ret;
    if (!std::filesystem::is_directory(folder))
      {
        if (output_messages)
          {
            std::cout << "ERROR: '" << folder << "' is not a valid folder." << std::endl;
          }
        return ret;
      }
    std::set<std::string> read_one_part_of_v1_cells;
    int event_count = 0;
    for (const std::filesystem::path & file : std::filesystem::directory_iterator(folder))
      {
        if ( max_events > 0 && event_count >= max_events &&
             ret.geometry.size() > 0 && ret.noise.size() > 0 )
          {
            break;
          }
 
        const bool can_load_events = (max_events < 0) || (event_count < max_events);
        const std::string filename = file.stem().native();
 
        if (filter_function(filename))
          {
            continue;
          }
 
        auto check_error = [&](const ErrorState & es, const std::string & str)
        {
          if (es == ErrorState::OK)
            {
              return false;
            }
          if (output_messages)
            {
              std::cout << "ERROR: '" << file << "' is not a valid " << str << " file (" << (int) es << ")." << std::endl;
            }
          return true;
        };
 
        auto output_loading_message = [&](const std::string & str)
        {
          if (output_messages)
            {
              std::cout << "Loaded " << str << " from '" << file << "'." << std::endl;
            }
        };
 
        if (file.extension() == ".geometry")
          {
            if (!flo.load_geometry)
              {
                continue;
              }
            CaloRecGPU::Helpers::CPU_object<CaloRecGPU::GeometryArr> tempgeo(true);
            if (check_error(ConstantInformation::read_geometry(file, tempgeo), "geometry"))
              {
                continue;
              }
            ret.geometry[filename] = std::move(tempgeo);
            output_loading_message("geometry");
          }
        else if (file.extension() == ".cellinfo")
          {
            if (!flo.load_cell_info || !can_load_events)
              {
                continue;
              }
            CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellInfoArr> tempcellinfo(true);
            if (check_error(EventInformation::read_cell_info(file, tempcellinfo), "cell info"))
              {
                continue;
              }
            ret.cell_info[filename] = std::move(tempcellinfo);
            output_loading_message("cell info");
            if (ret.cluster_info.count(filename) > 0 || !flo.load_cluster_info)
              {
                ++event_count;
              }
          }
        else if (file.extension() == ".clusterinfo")
          {
            if (!flo.load_cluster_info || !can_load_events)
              {
                continue;
              }
            CaloRecGPU::Helpers::CPU_object<CaloRecGPU::ClusterInfoArr> tempclu(true);
            if (check_error(EventInformation::read_cluster_info(file, tempclu), "cluster info"))
              {
                continue;
              }
            ret.cluster_info[filename] = std::move(tempclu);
            output_loading_message("cluster info");
            if (ret.cell_info.count(filename) > 0 || !flo.load_cell_info)
              {
                ++event_count;
              }
          }
        else if (file.extension() == ".noise")
          {
            if (!flo.load_noise)
              {
                continue;
              }
            CaloRecGPU::Helpers::CPU_object<CaloRecGPU::CellNoiseArr> tempnois(true);
            if (check_error(ConstantInformation::read_noise(file, tempnois), "noise"))
              {
                continue;
              }
            ret.noise[filename] = std::move(tempnois);
            output_loading_message("noise");
          }
      }
    return ret;
  }
 
  inline static FolderLoad load_folder(const std::filesystem::path & folder,
                                       int max_events = -1,
                                       const FolderLoadOptions & flo = FolderLoadOptions::None(),
                                       const bool output_messages = true)
  {
    return load_folder_filter([&](const std::string &)
    {
      return false;
    },
    folder, max_events, flo, output_messages);
  }
 
};
 
#endif //CALORECGPU_STANDALONEDATAIO_H