Two-stage MUSIC pairs: grid-boundary flag reset, MUSIC4GPU pin, droplet pruning

external_packages/get_music4gpu.sh

@@ -15,11 +15,17 @@
 
 # using a commit from the MUSIC repository that is compatible with the current X-SCAPE version
 folderName="music4gpu"
+# XSCAPE branch with the jet source slot (add_hydro_source_terms_from_jet, which
+# MusicWrapper calls since X-SCAPE PR #138) and per-step jet droplet pruning
+commitHash="3037be76518e4ae17b42c2236392bc93603c9f89"
 
 git clone https://github.com/jhputschke/MUSIC4GPU.git -b XSCAPE $folderName
 cd $folderName
+git checkout $commitHash
 cd EOS
 bash download_hotQCD.sh SMASH_binary
+# EOS 9 (UrQMD hadron list), used by e.g. the PyJetscape prod_AuAu_0_10 productions
+bash download_hotQCD.sh binary
 
 # Finite muB not supprted on GPU yet!!!
 # Download the 4D EoS tables (only needed for EOS 20, only download if necessary, large files)

src/framework/LiquefierBase.cc

@@ -31,7 +31,13 @@ void LiquefierBase::get_source(Jetscape::real tau, Jetscape::real x,
                                Jetscape::real y, Jetscape::real eta,
                                std::array<Jetscape::real, 4> &jmu) const {
   jmu = {0.0, 0.0, 0.0, 0.0};
-  for (const auto &drop_i : dropletlist) {
+  // Inside a prepared window only the droplets that can contribute there are
+  // visited, in their original order; the others would add exactly zero.
+  const bool pruned = (tau >= active_tau_lo_ && tau <= active_tau_hi_);
+  const std::size_t n_loop =
+      pruned ? active_droplets_.size() : dropletlist.size();
+  for (std::size_t k = 0; k < n_loop; k++) {
+    const auto &drop_i = dropletlist[pruned ? active_droplets_[k] : k];
     const auto x_drop = drop_i.get_xmu();
     double ds2 = tau * tau + x_drop[0] * x_drop[0] -
                  2.0 * tau * x_drop[0] * cosh(eta - x_drop[3]) -
@@ -227,7 +233,21 @@ void LiquefierBase::add_hydro_sources(std::vector<Parton> &pIn,
   }
 }
 
-void LiquefierBase::ClearTask() { dropletlist.clear(); }
+void LiquefierBase::ClearTask() {
+  dropletlist.clear();
+  invalidate_active_droplets();
+}
+
+void LiquefierBase::prepare_active_droplets(double tau_lo, double tau_hi) {
+  active_droplets_.clear();
+  for (std::size_t i = 0; i < dropletlist.size(); i++) {
+    if (droplet_may_contribute(dropletlist[i], tau_lo, tau_hi)) {
+      active_droplets_.push_back(static_cast<int>(i));
+    }
+  }
+  active_tau_lo_ = tau_lo;
+  active_tau_hi_ = tau_hi;
+}
 
 Jetscape::real LiquefierBase::get_dropletlist_total_energy() const {
   Jetscape::real total_E = 0.0;

src/framework/LiquefierBase.h

@@ -91,6 +91,13 @@ class LiquefierBase {
  private:
   std::vector<Droplet>
       dropletlist;  ///< List of droplets representing source contributions
+  /// Indices into dropletlist of the droplets that can contribute to a query
+  /// time in [active_tau_lo_, active_tau_hi_]; see prepare_active_droplets().
+  /// An empty window (lo > hi) means not prepared: get_source() then loops
+  /// over every droplet.
+  std::vector<int> active_droplets_;
+  double active_tau_lo_ = 1.;
+  double active_tau_hi_ = 0.;
   bool GetHydroCellSignalConnected;  ///< Flag for whether signal connection to
                                      ///< hydro exists
   const int drop_stat;               ///< Droplet statistics
@@ -117,7 +124,44 @@ class LiquefierBase {
    * @brief Add a droplet to the internal list.
    * @param droplet_in Droplet to add
    */
-  void add_a_droplet(Droplet droplet_in) { dropletlist.push_back(droplet_in); }
+  void add_a_droplet(Droplet droplet_in) {
+    dropletlist.push_back(droplet_in);
+    invalidate_active_droplets();
+  }
+
+  /**
+   * @brief Restrict get_source() to the droplets that can contribute to a
+   * query time in [tau_lo, tau_hi], e.g. one hydro time step with its
+   * Runge-Kutta substeps.
+   *
+   * Uses droplet_may_contribute(). Queries outside the window, and any query
+   * before the first call, still loop over every droplet, so a caller that
+   * never prepares gets the unpruned result. Adding a droplet or clearing
+   * the list drops the window.
+   */
+  void prepare_active_droplets(double tau_lo, double tau_hi);
+
+  /// Drop the prepared window: get_source() loops over every droplet again.
+  void invalidate_active_droplets() {
+    active_droplets_.clear();
+    active_tau_lo_ = 1.;
+    active_tau_hi_ = 0.;
+  }
+
+  /// Droplets kept by the last prepare_active_droplets() call.
+  int get_number_of_active_droplets() const {
+    return static_cast<int>(active_droplets_.size());
+  }
+
+  /**
+   * @brief Whether drop_i can give a non-zero smearing_kernel() for some
+   * query time in [tau_lo, tau_hi]. Must never return false for a droplet
+   * that contributes. The default keeps every droplet (no pruning).
+   */
+  virtual bool droplet_may_contribute(const Droplet &drop_i, double tau_lo,
+                                      double tau_hi) const {
+    return true;
+  }
 
   /**
    * @brief Get number of droplet conversions performed.

src/hydro/MusicWrapper.cc

@@ -438,12 +438,16 @@ void MpiMusic::EvolveHydroUpto(const double tauEnd) {
     }
     music_hydro_ptr->prepare_run_hydro_one_time_step();
     hydro_source_terms_ptr->set_source_tau_max(GetSourceTermTauMax());
+    // see EvolveHydro(): MUSIC never clears this flag itself
+    music_hydro_ptr->setReRunHydro(false);
+    hit_grid_boundary_ = false;
   }
 
   if (hydro_status != FINISHED) {
     int status = music_hydro_ptr->run_hydro_upto(tauEnd);
     if (status != 0) {
       hydro_status = FINISHED;
+      WarnIfGridBoundaryHit();
     }
   }
   // PassHydroSurfaceToFramework();
@@ -519,8 +523,16 @@ void MpiMusic::EvolveHydro() {
 
   if (hydro_status == INITIALIZED) {
     JSINFO << "running MUSIC ...";
+    // MUSIC stops when its freeze-out surface reaches the transverse grid
+    // edge and sets reRunHydro, expecting its driver to re-run on a larger
+    // grid (MUSIC's own main.cpp does). X-SCAPE does not re-run, and MUSIC
+    // clears the flag only in its constructor, so without this reset every
+    // later event of this instance would stop at its first freeze-out check.
+    music_hydro_ptr->setReRunHydro(false);
+    hit_grid_boundary_ = false;
     music_hydro_ptr->run_hydro();
     hydro_status = FINISHED;
+    WarnIfGridBoundaryHit();
   }
 
   if (dump_hydro_only_) {
@@ -583,6 +595,17 @@ void MpiMusic::EvolveHydro() {
   }
 }
 
+void MpiMusic::WarnIfGridBoundaryHit() {
+  hit_grid_boundary_ = music_hydro_ptr->getReRunHydro();
+  if (hit_grid_boundary_) {
+    JSWARN << "MUSIC " << GetId() << ": the freeze-out surface reached the "
+           << "transverse grid boundary, so MUSIC stopped this event's "
+           << "evolution there (it expects a re-run on a larger grid, which "
+           << "X-SCAPE does not do). The stored evolution is truncated; "
+           << "enlarge <IS><grid_max_x>/<grid_max_y>.";
+  }
+}
+
 void MpiMusic::Clear() {
   if (preserve_bulk_info_) {
     // Keep bulk_info.data alive so Python can inspect it after Exec().

src/hydro/MusicWrapper.h

@@ -37,6 +37,10 @@ class HydroSourceJETSCAPE : public HydroSourceBase {
 
   double dtau;
 
+  // Per-step droplet pruning, see prepare_list_for_current_tau_frame().
+  double pruning_step_ = 0.1;  // [fm]; the hydro dtau once set_hydro_dtau() ran
+  std::shared_ptr<LiquefierBase> liquefier_step_;  // locked once per step
+
  public:
   HydroSourceJETSCAPE() = default;
   ~HydroSourceJETSCAPE() {}
@@ -46,13 +50,31 @@ class HydroSourceJETSCAPE : public HydroSourceBase {
   // add the value to the hadronic liquefier
   void set_hydro_dtau(double val) {
     dtau = val;
+    pruning_step_ = val;
     if (!weak_ptr_is_uninitialized(hadronic_liquefier_ptr)) {
       hadronic_liquefier_ptr.lock()->set_hydro_dtau(dtau);
     }
   };
 
   void add_a_liquefier(std::shared_ptr<LiquefierBase> new_liquefier) {
     liquefier_ptr = new_liquefier;
+    liquefier_step_.reset();
+  }
+
+  //! Called by MUSIC once per time step, before its Runge-Kutta substeps
+  //! (query times tau and tau + dtau). A CausalLiquefier droplet deposits only
+  //! in the step containing tau_drop + tau_delay, so the liquefier keeps just
+  //! the droplets that can deposit at a query time in
+  //! [tau - step, tau + 2 step]; the result is unchanged, because the others
+  //! add exactly zero, and queries outside the window still see every droplet.
+  //! Also locks the liquefier once per step instead of once per cell.
+  void prepare_list_for_current_tau_frame(const double tau_local) {
+    if (weak_ptr_is_uninitialized(liquefier_ptr)) return;
+    liquefier_step_ = liquefier_ptr.lock();
+    if (liquefier_step_) {
+      liquefier_step_->prepare_active_droplets(tau_local - pruning_step_,
+                                               tau_local + 2. * pruning_step_);
+    }
   }
 
   void add_a_hadronic_liquefier(
@@ -134,7 +156,11 @@ class HydroSourceJETSCAPE : public HydroSourceBase {
     j_mu = {0.0};
     if (!weak_ptr_is_uninitialized(liquefier_ptr)) {
       std::array<Jetscape::real, 4> jmu_tmp = {0.0};
-      liquefier_ptr.lock()->get_source(tau, x, y, eta_s, jmu_tmp);
+      if (liquefier_step_) {
+        liquefier_step_->get_source(tau, x, y, eta_s, jmu_tmp);
+      } else {
+        liquefier_ptr.lock()->get_source(tau, x, y, eta_s, jmu_tmp);
+      }
       for (int i = 0; i < 4; i++) {
         j_mu[i] =
             jmu_tmp[i] / hbarC;  // convert the unit from GeV/fm^4 to 1/fm^5
@@ -274,6 +300,11 @@ class MpiMusic : public FluidDynamics {
   bool get_skip_surface() const { return skip_surface_; }
   void set_skip_surface(bool v) { skip_surface_ = v; }
 
+  // True if the last evolution stopped because the freeze-out surface
+  // reached the transverse grid boundary (MUSIC's reRunHydro): the stored
+  // evolution is then truncated. Reset at the start of every evolution.
+  bool get_hit_grid_boundary() const { return hit_grid_boundary_; }
+
   // Thin pass-throughs to MUSIC's native in-memory store (music_hydro_ptr is
   // private, so the writer reaches it through these). music_hydro_ptr is only
   // created in InitializeHydro(), so guard against calls before that (Python).
@@ -333,6 +364,8 @@ class MpiMusic : public FluidDynamics {
   // Skip the freeze-out surface hand-off / file collection (see
   // set_skip_surface); read from <Hydro><MUSIC><skip_surface>.
   bool skip_surface_ = false;
+  bool hit_grid_boundary_ = false;
+  void WarnIfGridBoundaryHit();
 };
 
 #endif  // MUSICWRAPPER_H

src/liquefier/CausalLiquefier.cc

@@ -231,6 +231,17 @@ double CausalLiquefier::get_peta(double p0, double p3, double eta) const {
   return p3 * cosh(eta) - p0 * sinh(eta);
 }
 
+// smearing_kernel() deposits a droplet only in the step that contains
+// tau_drop + tau_delay: tau - dtau/2 <= tau_drop + tau_delay < tau + dtau/2.
+// The extra 1e-3 fm covers tau reaching get_source() as a float.
+bool CausalLiquefier::droplet_may_contribute(const Droplet &drop_i,
+                                             double tau_lo,
+                                             double tau_hi) const {
+  const double tau_deposit = drop_i.get_xmu()[0] + tau_delay;
+  const double margin = 0.5 * dtau + 1e-3;
+  return tau_deposit >= tau_lo - margin && tau_deposit <= tau_hi + margin;
+}
+
 // For debug, Change tau_delay
 void CausalLiquefier::set_t_delay(double new_tau_delay) {
   tau_delay = new_tau_delay;

src/liquefier/CausalLiquefier.h

@@ -52,6 +52,9 @@ class CausalLiquefier : public Jetscape::LiquefierBase {
                        Jetscape::real eta, const Droplet drop_i,
                        std::array<Jetscape::real, 4> &jmu) const;
 
+  bool droplet_may_contribute(const Droplet &drop_i, double tau_lo,
+                              double tau_hi) const override;
+
   double dumping(double t) const;
 
   double kernel_rho(double t, double r) const;