outputformatter.h

Go to the documentation of this file.

/*
 *
 *    Copyright (c) 2024-2025
 *      SMASH Team
 *
 *    GNU General Public License (GPLv3 or later)
 *
 */
 
#ifndef SRC_INCLUDE_SMASH_OUTPUTFORMATTER_H_
#define SRC_INCLUDE_SMASH_OUTPUTFORMATTER_H_
 
#include <functional>
#include <map>
#include <optional>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
#include "smash/particledata.h"
#include "smash/particles.h"
 
namespace smash {
 
struct ToASCII {
  using type = std::string;
 
  static constexpr std::optional<char> separator{' '};
 
  static constexpr std::optional<char> end_of_line{'\n'};
 
  type as_integer(int value) const { return std::to_string(value); }
 
  type as_double(double value) {
    constexpr size_t kBufferSize = 13;
    char buffer[kBufferSize];
    const auto length = std::snprintf(buffer, kBufferSize, "%g", value);
    assert(static_cast<size_t>(length) < kBufferSize);
    assert(length > 0);
    return std::string{buffer, buffer + length};
  }
 
  type as_precise_double(double value) {
    constexpr size_t kBufferSize = 16;
    char buffer[kBufferSize];
    const auto length = std::snprintf(buffer, kBufferSize, "%.9g", value);
    assert(static_cast<size_t>(length) < kBufferSize);
    assert(length > 0);
    return std::string{buffer, buffer + length};
  }
 
  type as_string(const std::string& str) const { return str; }
};
 
class ToBinary {
 public:
  using type = std::vector<char>;
 
  static constexpr std::optional<char> separator = std::nullopt;
 
  static constexpr std::optional<char> end_of_line = std::nullopt;
 
  type as_integer(int32_t value) const { return as_binary_data(value); }
 
  type as_double(double value) const { return as_binary_data(value); }
 
  type as_precise_double(double value) const { return as_double(value); }
 
  type as_string(const std::string& str) const {
    return type(str.begin(), str.end());
  }
 
 private:
  template <typename T>
  type as_binary_data(T value) const {
    type binary_data(sizeof(T));
    std::memcpy(binary_data.data(), &value, sizeof(T));
    return binary_data;
  }
};
 
template <typename Converter,
          std::enable_if_t<std::is_same_v<Converter, ToASCII> ||
                               std::is_same_v<Converter, ToBinary>,
                           bool> = true>
class OutputFormatter {
 public:
  explicit OutputFormatter(const std::vector<std::string>& in_quantities)
      : quantities_(in_quantities) {
    validate_quantities();
    for (const std::string& quantity : quantities_) {
      if (quantity == "t") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.position()[0]);
        });
      } else if (quantity == "x") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.position()[1]);
        });
      } else if (quantity == "y") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.position()[2]);
        });
      } else if (quantity == "z") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.position()[3]);
        });
      } else if (quantity == "mass") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.effective_mass());
        });
      } else if (quantity == "p0") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_precise_double(in.momentum()[0]);
        });
      } else if (quantity == "px") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_precise_double(in.momentum()[1]);
        });
      } else if (quantity == "py") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_precise_double(in.momentum()[2]);
        });
      } else if (quantity == "pz") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_precise_double(in.momentum()[3]);
        });
      } else if (quantity == "pdg") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.pdgcode().get_decimal());
        });
      } else if (quantity == "ID" || quantity == "id") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.id());
        });
      } else if (quantity == "charge") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.type().charge());
        });
      } else if (quantity == "ncoll") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(
              in.get_history().collisions_per_particle);
        });
      } else if (quantity == "form_time") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.formation_time());
        });
      } else if (quantity == "xsecfac") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.xsec_scaling_factor());
        });
      } else if (quantity == "proc_id_origin") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.get_history().id_process);
        });
      } else if (quantity == "proc_type_origin") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(
              static_cast<int>(in.get_history().process_type));
        });
      } else if (quantity == "time_last_coll") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(
              in.get_history().time_last_collision);
        });
      } else if (quantity == "pdg_mother1") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.get_history().p1.get_decimal());
        });
      } else if (quantity == "pdg_mother2") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.get_history().p2.get_decimal());
        });
      } else if (quantity == "baryon_number") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.pdgcode().baryon_number());
        });
      } else if (quantity == "strangeness") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_integer(in.pdgcode().strangeness());
        });
      } else if (quantity == "spin0") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.spin_vector()[0]);
        });
      } else if (quantity == "spinx") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.spin_vector()[1]);
        });
      } else if (quantity == "spiny") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.spin_vector()[2]);
        });
      } else if (quantity == "spinz") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.spin_vector()[3]);
        });
      } else if (quantity == "0") {  // for OSCAR1999
        getters_.push_back([this]([[maybe_unused]] const ParticleData& in) {
          return this->converter_.as_integer(0);
        });
      } else if (quantity == "tau") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.hyperbolic_time());
        });
      } else if (quantity == "eta" || quantity == "eta_s") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.spatial_rapidity());
        });
      } else if (quantity == "mt") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.transverse_mass());
        });
      } else if (quantity == "Rap" || quantity == "y_rap") {
        // "Rap" is used for compatibility with vHLLE
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.rapidity());
        });
      } else if (quantity == "perturbative_weight") {
        getters_.push_back([this](const ParticleData& in) {
          return this->converter_.as_double(in.perturbative_weight());
        });
      }
    }
  }
 
  std::size_t compute_single_size(const ParticleData& sample) const {
    const std::size_t n = getters_.size();
    if (n == 0) {
      return 0;
    }
    std::size_t size = 0;
    for (const auto& getter : getters_) {
      const typename Converter::type tmp = getter(sample);
      size += tmp.size();
    }
 
    if (Converter::separator.has_value()) {
      size += (n - 1);
    }
 
    if (Converter::end_of_line.has_value()) {
      size += 1;
    }
 
    return size;
  }
  typename Converter::type single_particle_data(const ParticleData& p) const {
    typename Converter::type chunk{};
    append_to_buffer(p, chunk);
    return chunk;
  }
 
  template <class Range,
            std::enable_if_t<std::is_same_v<Range, Particles> ||
                                 std::is_same_v<Range, ParticleList>,
                             bool> = true>
  typename Converter::type particles_data_chunk(const Range& particles) const {
    typename Converter::type chunk{};
    if (particles.size() == 0)
      return chunk;
 
    chunk.reserve(particles.size() * compute_single_size(particles.front()));
 
    for (const ParticleData& p : particles) {
      append_to_buffer(p, chunk);
    }
    return chunk;
  }
 
  typename Converter::type quantities_line() const {
    typename Converter::type out{};
 
    // Educated size guess to reduce the number of reallocations in push_backs:
    out.reserve(quantities_.size() * 5);
 
    for (const auto& string_name : quantities_) {
      if constexpr (Converter::separator) {
        if (!out.empty()) {
          out.push_back(*Converter::separator);
        }
      }
      typename Converter::type name = converter_.as_string(string_name);
      out.insert(out.end(), name.begin(), name.end());
    }
    if constexpr (Converter::end_of_line) {
      out.push_back(*Converter::end_of_line);
    }
    return out;
  }
 
  typename Converter::type unit_line() const {
    typename Converter::type out{};
 
    // Educated size guess to reduce the number of reallocations in push_backs:
    out.reserve(quantities_.size() * 5);
 
    for (const auto& key : quantities_) {
      if constexpr (Converter::separator) {
        if (!out.empty()) {
          out.push_back(*Converter::separator);
        }
      }
      const auto& unit = converter_.as_string(units_.at(key));
      out.insert(out.end(), unit.begin(), unit.end());
    }
    if constexpr (Converter::end_of_line) {
      out.push_back(*Converter::end_of_line);
    }
    return out;
  }
 
  struct UnknownQuantity : public std::invalid_argument {
    using std::invalid_argument::invalid_argument;
  };
 
  struct RepeatedQuantity : public std::invalid_argument {
    using std::invalid_argument::invalid_argument;
  };
 
  struct AliasesQuantity : public std::invalid_argument {
    using std::invalid_argument::invalid_argument;
  };
 
 private:
  Converter converter_{};
 
  std::vector<std::string> quantities_{};
 
  std::vector<std::function<typename Converter::type(const ParticleData&)>>
      getters_{};
 
  const std::map<std::string, std::string> units_ = {
      {"t", "fm"},
      {"x", "fm"},
      {"y", "fm"},
      {"z", "fm"},
      {"mass", "GeV"},
      {"p0", "GeV"},
      {"px", "GeV"},
      {"py", "GeV"},
      {"pz", "GeV"},
      {"pdg", "none"},
      {"ID", "none"},
      {"id", "none"},  // used in OSCAR1999
      {"charge", "e"},
      {"ncoll", "none"},
      {"form_time", "fm"},
      {"xsecfac", "none"},
      {"proc_id_origin", "none"},
      {"proc_type_origin", "none"},
      {"time_last_coll", "fm"},
      {"pdg_mother1", "none"},
      {"pdg_mother2", "none"},
      {"baryon_number", "none"},
      {"strangeness", "none"},
      {"0", "0"},  // for OSCAR1999;
      {"tau", "fm"},
      {"eta", "none"},
      {"eta_s", "none"},
      {"mt", "GeV"},
      {"Rap", "none"},
      {"y_rap", "none"},
      {"spin0", "none"},
      {"spinx", "none"},
      {"spiny", "none"},
      {"spinz", "none"},
      {"perturbative_weight", "none"}};
 
  const std::vector<std::pair<std::string, std::string>> aliases_ = {
      {"ID", "id"}, {"eta_s", "eta"}, {"y_rap", "Rap"}};
 
  void validate_quantities() {
    if (quantities_.empty()) {
      throw std::invalid_argument(
          "OutputFormatter: Empty quantities handed over to the class.");
    }
    std::string repeated{};
    for (const std::string& quantity : quantities_) {
      if (std::count(quantities_.begin(), quantities_.end(), quantity) > 1) {
        repeated += "'" + quantity + "',";
      }
    }
    if (!repeated.empty()) {
      throw RepeatedQuantity("Repeated \"Quantities\": " + repeated +
                             " please fix the configuration file.\n");
    }
    std::string unknown{};
    for (const std::string& quantity : quantities_) {
      if (units_.count(quantity) == 0) {
        unknown += "'" + quantity + "',";
      }
    }
    if (!unknown.empty()) {
      throw UnknownQuantity("Unknown \"Quantities\": " + unknown +
                            " please fix the configuration file.\n");
    }
    for (const auto& alias : aliases_) {
      const bool first_quantity_is_given =
          std::find(quantities_.begin(), quantities_.end(), alias.first) !=
          quantities_.end();
      const bool second_quantity_is_given =
          std::find(quantities_.begin(), quantities_.end(), alias.second) !=
          quantities_.end();
      if (first_quantity_is_given && second_quantity_is_given) {
        throw AliasesQuantity(
            "Both '" + alias.first + "' and '" + alias.second +
            "' cannot be provided in the \"Quantities\" key together. Please, "
            "fix the configuration file.");
      }
    }
  }
 
  void append_to_buffer(const ParticleData& particle,
                        typename Converter::type& buffer) const {
    for (std::size_t i = 0; i < getters_.size(); ++i) {
      if constexpr (Converter::separator) {
        if (i > 0)
          buffer.push_back(*Converter::separator);
      }
      auto data = getters_[i](particle);
      buffer.insert(buffer.end(), std::make_move_iterator(data.begin()),
                    std::make_move_iterator(data.end()));
    }
    if constexpr (Converter::end_of_line) {
      buffer.push_back(*Converter::end_of_line);
    }
  }
};
 
namespace details {
 
template <typename Converter, class Range,
          std::enable_if_t<std::is_same_v<Range, Particles> ||
                               std::is_same_v<Range, ParticleList>,
                           bool> = true>
void write_in_chunk_impl(
    const Range& particles, const OutputFormatter<Converter>& formatter,
    std::function<void(const typename Converter::type&)> write,
    std::size_t max_buffer_bytes = 1'000'000'000) {
  if (particles.size() == 0)
    return;
 
  const std::size_t bytes_per_particle =
      formatter.compute_single_size(particles.front());
 
  if (2.0 * bytes_per_particle > max_buffer_bytes) {
    throw std::runtime_error(
        "write_in_chunk_impl: the estimated size of a single particle line "
        "exceeds half of the configured max_buffer_bytes.\n"
        "This effectively means only one particle would fit per chunk, "
        "which defeats the purpose of chunked writing.\n"
        "Increase max_buffer_bytes to at least twice the particle line size "
        "to use this function correctly.");
  }
 
  if (particles.size() * bytes_per_particle <= max_buffer_bytes) {
    write(formatter.particles_data_chunk(particles));
    return;
  }
 
  using Buffer = typename Converter::type;
  Buffer buffer;
  buffer.reserve(max_buffer_bytes);
  std::size_t current_size = 0;
 
  for (const auto& particle : particles) {
    Buffer line = formatter.single_particle_data(particle);
    const std::size_t line_size = line.size();
    if (current_size + line_size > max_buffer_bytes) {
      write(buffer);
      buffer.clear();
      current_size = 0;
    }
    buffer.insert(buffer.end(), std::make_move_iterator(line.begin()),
                  std::make_move_iterator(line.end()));
    current_size += line_size;
  }
 
  if (!buffer.empty()) {
    write(buffer);
  }
}
}  // namespace details
 
template <typename Converter, class Range,
          std::enable_if_t<std::is_same_v<Range, Particles> ||
                               std::is_same_v<Range, ParticleList>,
                           bool> = true>
void write_in_chunk(
    const Range& particles, const OutputFormatter<Converter>& formatter,
    std::function<void(const typename Converter::type&)> write) {
  details::write_in_chunk_impl(particles, formatter, write);
}
 
}  // namespace smash
 
#endif  // SRC_INCLUDE_SMASH_OUTPUTFORMATTER_H_