rootoutput.cc

Go to the documentation of this file.

/*
 *
 *    Copyright (c) 2014-2024
 *      SMASH Team
 *
 *    GNU General Public License (GPLv3 or later)
 *
 */
 
#include "smash/rootoutput.h"
 
#include "TFile.h"
#include "TTree.h"
#include "smash/action.h"
#include "smash/clock.h"
#include "smash/forwarddeclarations.h"
#include "smash/particles.h"
 
namespace smash {
static constexpr int LHyperSurfaceCrossing = LogArea::HyperSurfaceCrossing::id;
 
const int RootOutput::max_buffer_size_ = 500000;
 
RootOutput::RootOutput(const std::filesystem::path &path,
                       const std::string &name, const OutputParameters &out_par)
    : OutputInterface(name),
      filename_(path / (name + ".root")),
      write_collisions_(name == "Collisions" || name == "Dileptons" ||
                        name == "Photons"),
      write_particles_(name == "Particles"),
      write_initial_conditions_(name == "SMASH_IC"),
      particles_only_final_(out_par.part_only_final),
      part_extended_(out_par.part_extended),
      coll_extended_(out_par.coll_extended),
      ic_extended_(out_par.ic_extended) {
  filename_unfinished_ = filename_;
  filename_unfinished_ += ".unfinished";
  root_out_file_ =
      std::make_unique<TFile>(filename_unfinished_.native().c_str(), "NEW");
  init_trees();
}
 
void RootOutput::init_trees() {
  if (write_particles_ || write_initial_conditions_) {
    particles_tree_ = new TTree("particles", "particles");
 
    particles_tree_->Branch("ev", &ev_, "ev/I");
    particles_tree_->Branch("ens", &ens_, "ens/I");
    particles_tree_->Branch("tcounter", &tcounter_, "tcounter/I");
    particles_tree_->Branch("npart", &npart_, "npart/I");
    particles_tree_->Branch("test_p", &test_p_, "test_p/I");
    particles_tree_->Branch("modus_l", &modus_l_, "modus_l/D");
    particles_tree_->Branch("current_t", &current_t_, "current_t/D");
    particles_tree_->Branch("impact_b", &impact_b_, "impact_b/D");
    particles_tree_->Branch("empty_event", &empty_event_, "empty_event/O");
 
    particles_tree_->Branch("pdgcode", &pdgcode_[0], "pdgcode[npart]/I");
    particles_tree_->Branch("charge", &charge_[0], "charge[npart]/I");
 
    particles_tree_->Branch("p0", &p0_[0], "p0[npart]/D");
    particles_tree_->Branch("px", &px_[0], "px[npart]/D");
    particles_tree_->Branch("py", &py_[0], "py[npart]/D");
    particles_tree_->Branch("pz", &pz_[0], "pz[npart]/D");
 
    particles_tree_->Branch("t", &t_[0], "t[npart]/D");
    particles_tree_->Branch("x", &x_[0], "x[npart]/D");
    particles_tree_->Branch("y", &y_[0], "y[npart]/D");
    particles_tree_->Branch("z", &z_[0], "z[npart]/D");
 
    particles_tree_->Branch("E_kinetic_tot", &E_kinetic_tot_,
                            "E_kinetic_tot/D");
    particles_tree_->Branch("E_fields_tot", &E_fields_tot_, "E_fields_tot/D");
    particles_tree_->Branch("E_tot", &E_tot_, "E_tot/D");
 
    if (part_extended_ || ic_extended_) {
      particles_tree_->Branch("ncoll", &coll_per_part_[0], "ncoll[npart]/I");
      particles_tree_->Branch("form_time", &formation_time_[0],
                              "form_time[npart]/D");
      particles_tree_->Branch("xsecfac", &xsec_factor_[0], "xsecfac[npart]/D");
      particles_tree_->Branch("proc_id_origin", &proc_id_origin_[0],
                              "proc_id_origin[npart]/I");
      particles_tree_->Branch("proc_type_origin", &proc_type_origin_[0],
                              "proc_type_origin[npart]/I");
      particles_tree_->Branch("time_last_coll", &time_last_coll_[0],
                              "time_last_coll[npart]/D");
      particles_tree_->Branch("pdg_mother1", &pdg_mother1_[0],
                              "pdg_mother1[npart]/I");
      particles_tree_->Branch("pdg_mother2", &pdg_mother2_[0],
                              "pdg_mother2[npart]/I");
      particles_tree_->Branch("baryon_number", &baryon_number_[0],
                              "baryon_number[npart]/I");
      particles_tree_->Branch("strangeness", &strangeness_[0],
                              "strangeness[npart]/I");
    }
  }
 
  if (write_collisions_) {
    collisions_tree_ = new TTree("collisions", "collisions");
 
    collisions_tree_->Branch("nin", &nin_, "nin/I");
    collisions_tree_->Branch("nout", &nout_, "nout/I");
    collisions_tree_->Branch("npart", &npart_, "npart/I");
    collisions_tree_->Branch("ev", &ev_, "ev/I");
    collisions_tree_->Branch("ens", &ens_, "ens/I");
    collisions_tree_->Branch("weight", &wgt_, "weight/D");
    collisions_tree_->Branch("partial_weight", &par_wgt_, "partial_weight/D");
 
    collisions_tree_->Branch("pdgcode", &pdgcode_[0], "pdgcode[npart]/I");
    collisions_tree_->Branch("charge", &charge_[0], "charge[npart]/I");
 
    collisions_tree_->Branch("p0", &p0_[0], "p0[npart]/D");
    collisions_tree_->Branch("px", &px_[0], "px[npart]/D");
    collisions_tree_->Branch("py", &py_[0], "py[npart]/D");
    collisions_tree_->Branch("pz", &pz_[0], "pz[npart]/D");
 
    collisions_tree_->Branch("t", &t_[0], "t[npart]/D");
    collisions_tree_->Branch("x", &x_[0], "x[npart]/D");
    collisions_tree_->Branch("y", &y_[0], "y[npart]/D");
    collisions_tree_->Branch("z", &z_[0], "z[npart]/D");
 
    if (coll_extended_) {
      collisions_tree_->Branch("ncoll", &coll_per_part_[0], "ncoll[npart]/I");
      collisions_tree_->Branch("form_time", &formation_time_[0],
                               "form_time[npart]/D");
      collisions_tree_->Branch("xsecfac", &xsec_factor_[0], "xsecfac[npart]/D");
      collisions_tree_->Branch("proc_id_origin", &proc_id_origin_[0],
                               "proc_id_origin[npart]/I");
      collisions_tree_->Branch("proc_type_origin", &proc_type_origin_[0],
                               "proc_type_origin[npart]/I");
      collisions_tree_->Branch("time_last_coll", &time_last_coll_[0],
                               "time_last_coll[npart]/D");
      collisions_tree_->Branch("pdg_mother1", &pdg_mother1_[0],
                               "pdg_mother1[npart]/I");
      collisions_tree_->Branch("pdg_mother2", &pdg_mother2_[0],
                               "pdg_mother2[npart]/I");
      collisions_tree_->Branch("baryon_number", &baryon_number_[0],
                               "baryon_number[npart]/I");
      collisions_tree_->Branch("strangeness", &strangeness_[0],
                               "strangeness[npart]/I");
    }
  }
}
 
RootOutput::~RootOutput() {
  // kOverwrite option prevents from writing extra TKey objects into root file
  root_out_file_->Write("", TObject::kOverwrite);
  root_out_file_->Close();
  std::filesystem::rename(filename_unfinished_, filename_);
}
 
void RootOutput::at_eventstart(const Particles &particles,
                               const EventLabel &event_label,
                               const EventInfo &event) {
  // save event and ensemble numbers
  current_event_ = event_label.event_number;
  current_ensemble_ = event_label.ensemble_number;
 
  modus_l_ = event.modus_length;
  test_p_ = event.test_particles;
  current_t_ = event.current_time;
  E_kinetic_tot_ = event.total_kinetic_energy;
  E_fields_tot_ = event.total_mean_field_energy;
  E_tot_ = event.total_energy;
 
  if (write_particles_ && particles_only_final_ == OutputOnlyFinal::No) {
    output_counter_ = 0;
    // This is to have only one output of positive impact parameter per event
    impact_b_ = -1.0;
    empty_event_ = false;
    particles_to_tree(particles);
    output_counter_++;
  }
}
 
void RootOutput::at_intermediate_time(const Particles &particles,
                                      const std::unique_ptr<Clock> &,
                                      const DensityParameters &,
                                      const EventLabel &event_label,
                                      const EventInfo &event) {
  current_ensemble_ = event_label.ensemble_number;
  modus_l_ = event.modus_length;
  test_p_ = event.test_particles;
  current_t_ = event.current_time;
  E_kinetic_tot_ = event.total_kinetic_energy;
  E_fields_tot_ = event.total_mean_field_energy;
  E_tot_ = event.total_energy;
 
  if (write_particles_ && particles_only_final_ == OutputOnlyFinal::No) {
    particles_to_tree(particles);
    output_counter_++;
  }
}
 
void RootOutput::at_eventend(const Particles &particles,
                             const EventLabel &event_label,
                             const EventInfo &event) {
  current_ensemble_ = event_label.ensemble_number;
  modus_l_ = event.modus_length;
  test_p_ = event.test_particles;
  current_t_ = event.current_time;
  E_kinetic_tot_ = event.total_kinetic_energy;
  E_fields_tot_ = event.total_mean_field_energy;
  E_tot_ = event.total_energy;
 
  impact_b_ = event.impact_parameter;
  empty_event_ = event.empty_event;
  if (write_particles_ &&
      !(event.empty_event &&
        particles_only_final_ == OutputOnlyFinal::IfNotEmpty)) {
    particles_to_tree(particles);
  }
  /* Calculate the autosave frequency to be used. In case multiple ensembles are
   * used this has to be at least 1, so set it to 1 if the number of ensembles
   * is larger than the initial value of the autosave frequency. The evaluation
   * of the actual frequency to be used has to be done only once, i.e. when the
   * autosave_frequency_ has still its meaningless initial negative value. */
  if (autosave_frequency_ < 0) {
    autosave_frequency_ = 1000 / event.n_ensembles;
    if (autosave_frequency_ == 0) {
      autosave_frequency_ = 1;
    }
  }
  /* Forced regular dump from operational memory to disk. Very demanding!
   * If program crashes written data will NOT be lost. */
  if (current_event_ > 0 && current_event_ % autosave_frequency_ == 0) {
    if (write_particles_ || write_initial_conditions_) {
      particles_tree_->AutoSave("SaveSelf");
    }
    if (write_collisions_) {
      collisions_tree_->AutoSave("SaveSelf");
    }
  }
 
  if (write_initial_conditions_) {
    // If the runtime is too short some particles might not yet have
    // reached the hypersurface. Warning is printed.
    if (particles.size() != 0 && !event.impose_kinematic_cut_for_SMASH_IC) {
      logg[LHyperSurfaceCrossing].warn(
          "End time might be too small for initial conditions output. "
          "Hypersurface has not yet been crossed by ",
          particles.size(), " particle(s).");
    }
  }
}
 
void RootOutput::at_interaction(const Action &action,
                                const double /*density*/) {
  if (write_collisions_) {
    collisions_to_tree(action.incoming_particles(), action.outgoing_particles(),
                       action.get_total_weight(), action.get_partial_weight());
  }
 
  if (write_initial_conditions_ &&
      action.get_type() == ProcessType::HyperSurfaceCrossing) {
    particles_to_tree(action.incoming_particles());
  }
}
 
template <typename T>
void RootOutput::particles_to_tree(T &particles) {
  int i = 0;
 
  ev_ = current_event_;
  ens_ = current_ensemble_;
  tcounter_ = output_counter_;
  bool exceeded_buffer_message = true;
 
  for (const auto &p : particles) {
    // Buffer full - flush to tree, else fill with particles
    if (i >= max_buffer_size_) {
      if (exceeded_buffer_message) {
        logg[LOutput].warn()
            << "\nThe number of particles N = " << particles.size()
            << " exceeds the maximum buffer size B = " << max_buffer_size_
            << ".\nceil(N/B) = "
            << std::ceil(particles.size() /
                         static_cast<double>(max_buffer_size_))
            << " separate ROOT Tree entries will be created at this output."
            << "\nMaximum buffer size (max_buffer_size_) can be changed in "
            << "rootoutput.h\n\n";
        exceeded_buffer_message = false;
      }
      npart_ = max_buffer_size_;
      i = 0;
      particles_tree_->Fill();
    } else {
      pdgcode_[i] = p.pdgcode().get_decimal();
      charge_[i] = p.type().charge();
 
      p0_[i] = p.momentum().x0();
      px_[i] = p.momentum().x1();
      py_[i] = p.momentum().x2();
      pz_[i] = p.momentum().x3();
 
      t_[i] = p.position().x0();
      x_[i] = p.position().x1();
      y_[i] = p.position().x2();
      z_[i] = p.position().x3();
 
      if (part_extended_ || ic_extended_) {
        const auto h = p.get_history();
        formation_time_[i] = p.formation_time();
        xsec_factor_[i] = p.xsec_scaling_factor();
        time_last_coll_[i] = h.time_last_collision;
        coll_per_part_[i] = h.collisions_per_particle;
        proc_id_origin_[i] = h.id_process;
        proc_type_origin_[i] = static_cast<int>(h.process_type);
        pdg_mother1_[i] = h.p1.get_decimal();
        pdg_mother2_[i] = h.p2.get_decimal();
        baryon_number_[i] = p.type().baryon_number();
        strangeness_[i] = p.type().strangeness();
      }
 
      i++;
    }
  }
  // Flush rest to tree
  if (i > 0) {
    npart_ = i;
    particles_tree_->Fill();
  }
}
 
void RootOutput::collisions_to_tree(const ParticleList &incoming,
                                    const ParticleList &outgoing,
                                    const double weight,
                                    const double partial_weight) {
  ev_ = current_event_;
  ens_ = current_ensemble_;
  nin_ = incoming.size();
  nout_ = outgoing.size();
  npart_ = nin_ + nout_;
  wgt_ = weight;
  par_wgt_ = partial_weight;
 
  int i = 0;
 
  /* It is assumed that nin + nout < max_buffer_size_
   * This is true for any possible reaction for current buffer size
   * But if one wants initial/final particles written to collisions
   * then implementation should be updated. */
 
  for (const ParticleList &plist : {incoming, outgoing}) {
    for (const auto &p : plist) {
      pdgcode_[i] = p.pdgcode().get_decimal();
      charge_[i] = p.type().charge();
 
      p0_[i] = p.momentum().x0();
      px_[i] = p.momentum().x1();
      py_[i] = p.momentum().x2();
      pz_[i] = p.momentum().x3();
 
      t_[i] = p.position().x0();
      x_[i] = p.position().x1();
      y_[i] = p.position().x2();
      z_[i] = p.position().x3();
 
      if (coll_extended_) {
        const auto h = p.get_history();
        formation_time_[i] = p.formation_time();
        xsec_factor_[i] = p.xsec_scaling_factor();
        time_last_coll_[i] = h.time_last_collision;
        coll_per_part_[i] = h.collisions_per_particle;
        proc_id_origin_[i] = h.id_process;
        proc_type_origin_[i] = static_cast<int>(h.process_type);
        pdg_mother1_[i] = h.p1.get_decimal();
        pdg_mother2_[i] = h.p2.get_decimal();
        baryon_number_[i] = p.type().baryon_number();
      }
 
      i++;
    }
  }
 
  collisions_tree_->Fill();
}
}  // namespace smash