Change output format to dump to summary

corsika/detail/modules/writers/InteractionWriter.inl

@@ -18,8 +18,7 @@ namespace corsika {
       : obsPlane_(obsPlane)
       , showerAxis_(axis)
       , interactionCounter_(0)
-      , showerId_(0)
-      , nSecondaries_(0) {}
+      , showerId_(0) {}
 
   template <typename TTracking, typename TOutput>
   template <typename TStackView>
@@ -32,47 +31,70 @@ namespace corsika {
     auto primary = vS.getProjectile();
     auto dX = showerAxis_.getProjectedX(primary.getPosition());
     CORSIKA_LOG_INFO("First interaction at dX {}", dX);
-    CORSIKA_LOG_INFO("Primary: {}, E {}", primary.getPID(), primary.getKineticEnergy());
-
-    // Get the displacement of the primary particle w.r.t. observation plane
-    Vector const dr_prim = primary.getPosition() - obsPlane_.getPlane().getCenter();
-    auto const p_prim = primary.getMomentum();
-
-    *(output_.getWriter())
-        << showerId_ << true << static_cast<int>(get_PDG(primary.getPID()))
-        << static_cast<float>(dr_prim.dot(obsPlane_.getXAxis()) / 1_m)
-        << static_cast<float>(dr_prim.dot(obsPlane_.getYAxis()) / 1_m)
-        << static_cast<float>(dr_prim.dot(obsPlane_.getPlane().getNormal()) / 1_m)
-        << static_cast<float>(p_prim.dot(obsPlane_.getXAxis()) / 1_GeV)
-        << static_cast<float>(p_prim.dot(obsPlane_.getYAxis()) / 1_GeV)
-        << static_cast<float>(p_prim.dot(obsPlane_.getPlane().getNormal()) / 1_GeV)
-        << static_cast<double>(primary.getTime() / 1_ns)
-        << static_cast<float>(dX / (1_g / 1_cm / 1_cm)) << parquet::EndRow;
+    CORSIKA_LOG_INFO("Primary: {}, E_kin {}", primary.getPID(),
+                     primary.getKineticEnergy());
+
+    // get the location of the primary w.r.t. observation plane
+    Vector const displacement = primary.getPosition() - obsPlane_.getPlane().getCenter();
+
+    auto const x = displacement.dot(obsPlane_.getXAxis());
+    auto const y = displacement.dot(obsPlane_.getYAxis());
+    auto const z = displacement.dot(obsPlane_.getPlane().getNormal());
+    auto const nx = primary.getDirection().dot(obsPlane_.getXAxis());
+    auto const ny = primary.getDirection().dot(obsPlane_.getYAxis());
+    auto const nz = primary.getDirection().dot(obsPlane_.getPlane().getNormal());
+
+    auto const px = primary.getMomentum().dot(obsPlane_.getXAxis());
+    auto const py = primary.getMomentum().dot(obsPlane_.getYAxis());
+    auto const pz = primary.getMomentum().dot(obsPlane_.getPlane().getNormal());
+
+    summary_["shower_" + std::to_string(showerId_)]["pdg"] =
+        static_cast<int>(get_PDG(primary.getPID()));
+    summary_["shower_" + std::to_string(showerId_)]["name"] =
+        static_cast<std::string>(get_name(primary.getPID()));
+    summary_["shower_" + std::to_string(showerId_)]["total_energy"] =
+        (primary.getKineticEnergy() + get_mass(primary.getPID())) / 1_GeV;
+    summary_["shower_" + std::to_string(showerId_)]["kinetic_energy"] =
+        primary.getKineticEnergy() / 1_GeV;
+    summary_["shower_" + std::to_string(showerId_)]["x"] = x / 1_m;
+    summary_["shower_" + std::to_string(showerId_)]["y"] = y / 1_m;
+    summary_["shower_" + std::to_string(showerId_)]["z"] = z / 1_m;
+    summary_["shower_" + std::to_string(showerId_)]["nx"] = static_cast<double>(nx);
+    summary_["shower_" + std::to_string(showerId_)]["ny"] = static_cast<double>(ny);
+    summary_["shower_" + std::to_string(showerId_)]["nz"] = static_cast<double>(nz);
+    summary_["shower_" + std::to_string(showerId_)]["px"] =
+        static_cast<double>(px / 1_GeV);
+    summary_["shower_" + std::to_string(showerId_)]["py"] =
+        static_cast<double>(py / 1_GeV);
+    summary_["shower_" + std::to_string(showerId_)]["pz"] =
+        static_cast<double>(pz / 1_GeV);
+    summary_["shower_" + std::to_string(showerId_)]["time"] = primary.getTime() / 1_s;
+    summary_["shower_" + std::to_string(showerId_)]["slant_depth"] =
+        dX / (1_g / 1_cm / 1_cm);
+
+    uint nSecondaries = 0;
 
     // Loop through secondaries
     auto particle = vS.begin();
     while (particle != vS.end()) {
 
-      // Get the displacement of the secondary w.r.t. observation plane
-      Vector const dr_2nd = particle.getPosition() - obsPlane_.getPlane().getCenter();
+      // Get the momentum of the secondary, write w.r.t. observation plane
       auto const p_2nd = particle.getMomentum();
 
       *(output_.getWriter())
-          << showerId_ << false << static_cast<int>(get_PDG(particle.getPID()))
-          << static_cast<float>(dr_2nd.dot(obsPlane_.getXAxis()) / 1_m)
-          << static_cast<float>(dr_2nd.dot(obsPlane_.getYAxis()) / 1_m)
-          << static_cast<float>(dr_2nd.dot(obsPlane_.getPlane().getNormal()) / 1_m)
+          << showerId_ << static_cast<int>(get_PDG(particle.getPID()))
           << static_cast<float>(p_2nd.dot(obsPlane_.getXAxis()) / 1_GeV)
           << static_cast<float>(p_2nd.dot(obsPlane_.getYAxis()) / 1_GeV)
           << static_cast<float>(p_2nd.dot(obsPlane_.getPlane().getNormal()) / 1_GeV)
-          << static_cast<double>(particle.getTime() / 1_ns)
-          << static_cast<float>(dX / (1_g / 1_cm / 1_cm)) << parquet::EndRow;
+          << parquet::EndRow;
 
       CORSIKA_LOG_INFO(" 2ndary: {}, E_kin {}", particle.getPID(),
                        particle.getKineticEnergy());
       ++particle;
-      ++nSecondaries_;
+      ++nSecondaries;
     }
+
+    summary_["shower_" + std::to_string(showerId_)]["n_secondaries"] = nSecondaries;
   }
 
   template <typename TTracking, typename TOutput>
@@ -80,32 +102,19 @@ namespace corsika {
       boost::filesystem::path const& directory) {
     output_.initStreamer((directory / ("interactions.parquet")).string());
 
-    output_.addField("primary", parquet::Repetition::REQUIRED, parquet::Type::BOOLEAN,
-                     parquet::ConvertedType::NONE);
     output_.addField("pdg", parquet::Repetition::REQUIRED, parquet::Type::INT32,
                      parquet::ConvertedType::INT_32);
-    output_.addField("x", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
-                     parquet::ConvertedType::NONE);
-    output_.addField("y", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
-                     parquet::ConvertedType::NONE);
-    output_.addField("z", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
-                     parquet::ConvertedType::NONE);
     output_.addField("px", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
                      parquet::ConvertedType::NONE);
     output_.addField("py", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
                      parquet::ConvertedType::NONE);
     output_.addField("pz", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
                      parquet::ConvertedType::NONE);
-    output_.addField("time", parquet::Repetition::REQUIRED, parquet::Type::DOUBLE,
-                     parquet::ConvertedType::NONE);
-    output_.addField("dX", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
-                     parquet::ConvertedType::NONE);
 
     output_.buildStreamer();
 
     showerId_ = 0;
     interactionCounter_ = 0;
-    nSecondaries_ = 0;
     summary_ = YAML::Node();
   }
 
@@ -114,13 +123,10 @@ namespace corsika {
       unsigned int const showerId) {
     showerId_ = showerId;
     interactionCounter_ = 0;
-    nSecondaries_ = 0;
   }
 
   template <typename TTracking, typename TOutput>
-  inline void InteractionWriter<TTracking, TOutput>::endOfShower(unsigned int const) {
-    summary_["shower_" + std::to_string(showerId_)]["n_secondaries"] = nSecondaries_;
-  }
+  inline void InteractionWriter<TTracking, TOutput>::endOfShower(unsigned int const) {}
 
   template <typename TTracking, typename TOutput>
   inline void InteractionWriter<TTracking, TOutput>::endOfLibrary() {

corsika/modules/writers/InteractionWriter.hpp

@@ -57,7 +57,6 @@ namespace corsika {
 
     unsigned int interactionCounter_;
     unsigned int showerId_;
-    unsigned int nSecondaries_;
 
     ParquetStreamer output_;
     YAML::Node summary_;

corsika/output/OutputManager.hpp

@@ -35,8 +35,10 @@ namespace corsika {
     /**
      * Construct an OutputManager instance with a name in a given directory.
      *
-     * @param name    The name of this output collection.
-     * @param dir     The directory where the output directory will be stored.
+     * @param name        The name of this output collection.
+     * @param vseed       The seed for the simulation (written to summary file)
+     * @param input_args  The command line arguments at runtime (written to summary file)
+     * @param dir         The directory where the output directory will be stored.
      */
     OutputManager(std::string const& name, const long& vseed,
                   std::string const& input_args, boost::filesystem::path const& dir);

python/corsika/io/outputs/interaction.py

@@ -9,15 +9,45 @@
 """
 
 import logging
+import pyarrow.parquet as pq
 import os.path as op
 from typing import Any
+import yaml
 
-import pyarrow.parquet as pq
 
 from ..converters import arrow_to_numpy
 from .output import Output
+from .primary import Particle
+
+class FirstInteraction(Particle):
+    """
+    Simple class that holds the information of the first interaction
+    This is a complete wrapper of the `Particle` class at the moment
+    but is "new" so that it can be distinguished using `type(thing)`
+    """
+    def __init__(self, prop_dict: dict):
+        """
+        Parameters
+        ----------
+        prop_dict: dict
+            the entry from the output yaml file for a single shower
+        """
+        Particle.__init__(self, prop_dict)
+
+    def __repr__(self) -> str:
+        """
+        Return a string representation of this class.
+        """
+        out_str = "FirstInteraction:\n"
+        for key in self.__dict__.keys():
+            out_str += f"\t{key}: {self.__dict__[key]}\n"
+        return out_str
 
 
+    @property
+    def momentum(self) -> list:
+        return [self.px, self.py, self.pz]
+
 class Interactions(Output):
     """
     Reads the interactions and secondaries
@@ -37,11 +67,18 @@ class Interactions(Output):
         # try and load our data
         try:
             self.__data = pq.read_table(op.join(path, "interactions.parquet"))
+            with open(op.join(path, "summary.yaml"), "r") as f:
+                temp = yaml.load(f, Loader=yaml.Loader)
+                self.__projectiles = [FirstInteraction(x) for x in temp.values()]
         except Exception as e:
             logging.getLogger("corsika").warn(
                 f"An error occured loading an Interaction: {e}"
             )
 
+    @property
+    def projectiles(self) -> list:
+        return self.__projectiles
+
     def is_good(self) -> bool:
         """
         Returns true if this output has been read successfully

python/examples/first_interactions.py

@@ -32,40 +32,50 @@ lib = corsika.Library(args.input_dir)
 # Load the primary particle information from the shower
 interactions = lib.get("interactions").data
 interaction_config = lib.get("interactions").config
+projectiles = lib.get("interactions").projectiles
 
 shower_ids = np.unique(interactions["shower"])
 
-ncols = 1
+ncols = 2
 nrows = len(shower_ids)
 fig, ax = plt.subplots(
     ncols=ncols,
     nrows=nrows,
-    figsize=(ncols * 8, nrows * 3),
+    figsize=(ncols * 8, nrows * 8),
     gridspec_kw={"wspace": 0.2, "hspace": 0.3},
 )
-ax = np.atleast_1d(ax).flatten()
+ax = np.atleast_2d(ax)
+if ax.shape == (1, 2):
+    ax = ax.T
 
-for sh_id in shower_ids:
-    data = interactions[interactions["shower"] == sh_id]
+for ish, sh_id in enumerate(shower_ids):
 
-    mother = data[data["primary"]]
-    norm = np.sqrt(mother["px"] ** 2 + mother["py"] ** 2 + mother["pz"] ** 2)
-    angle = np.arctan2(mother["py"].iloc[0], mother["px"].iloc[0])
-    pz = mother["pz"].iloc[0]
-    px = np.cos(-angle) * mother["px"].iloc[0] + np.sin(-angle) * mother["py"].iloc[0]
+    # Find all of the secondary particles associated with this shower
+    daughters = interactions[interactions["shower"] == sh_id]
 
-    pdg = int(mother["pdg"].iloc[0])
-    name = f"${particle.Particle.from_pdgid(pdg).latex_name}$"
-    ax[sh_id].plot([-px, 0], [-pz, 0], label=name)
+    mother = projectiles[ish]
+    norm = np.sqrt(sum(np.array(mother.momentum)**2)) # total momentum
+
+    print("Plotting shower", sh_id)
+    print(mother)
+
+    px, py, pz = mother.momentum
+
+    mother_name = f"${particle.Particle.from_pdgid(mother.pdg).latex_name}$"
+    ax[0, sh_id].plot([-px, 0], [-pz, 0], label=mother_name)
+    ax[1, sh_id].plot([-py, 0], [-pz, 0], label=mother_name)
 
     # Get daughters, filter out zero-momentum particles
-    daughters = data[data["primary"] is False]
     p_tot = np.sqrt(daughters["px"] ** 2 + daughters["py"] ** 2 + daughters["pz"] ** 2)
     daughters = daughters[p_tot > 0]
 
+    max_pz = abs(pz)
+    max_pxy = max(abs(px), abs(py))
+
     for i in range(len(daughters)):
         pz = daughters["pz"].iloc[i]
-        pt = np.sqrt(daughters["px"].iloc[i] ** 2 + daughters["py"].iloc[i] ** 2)
+        px = daughters["px"].iloc[i]
+        py = daughters["py"].iloc[i]
 
         pdg = int(daughters["pdg"].iloc[i])
         try:
@@ -73,16 +83,29 @@ for sh_id in shower_ids:
         except Exception:
             name = "Unknown"
 
-        ax[sh_id].plot([pt, 0], [pz, 0], label=name)
+        ax[0, sh_id].plot([px, 0], [pz, 0], label=name)
+        ax[1, sh_id].plot([py, 0], [pz, 0], label=name)
+
+        max_pz = max(max_pz, abs(pz))
+        max_pxy = max(max_pxy, max(abs(px), abs(py)))
+
+
+    for i in range(2):
+        ax[i, sh_id].legend()
+
+        ymin, ymax = ax[i, sh_id].get_ylim()
+        ax[i, sh_id].set_ylim(-max_pz, max_pz)
+        ax[i, sh_id].set_xlim(-max_pxy, max_pxy)
 
-    ax[sh_id].set_aspect("equal")
-    ax[sh_id].legend()
+        unit_energy_str = interaction_config["units"]["energy"]
+        title = f"Mother: {mother.kinetic_energy:0.2e} {unit_energy_str}"
+        title += ", " + mother_name
+        title += f"\n{mother.n_secondaries} secondaries"
+        ax[i, sh_id].set_title(title)
+        ax[i, sh_id].set_ylabel(f"$P_z$ ({unit_energy_str})")
 
-    ymin, ymax = ax[sh_id].get_ylim()
-    xmin, xmax = ax[sh_id].get_xlim()
-    maxVal = max(ymin, ymax, xmin, xmax)
-    ax[sh_id].set_ylim(-maxVal, maxVal)
-    ax[sh_id].set_xlim(-maxVal, maxVal)
+    ax[0, sh_id].set_xlabel(f"$P_x$ ({unit_energy_str})")
+    ax[1, sh_id].set_xlabel(f"$P_y$ ({unit_energy_str})")
 
 plot_path = os.path.join(args.output_dir, "interactions.png")
 print("Saving", plot_path)