From d2777471cc68bff382eac41ddd88f59d5e5a0b40 Mon Sep 17 00:00:00 2001
From: Alan Coleman <alanc@udel.edu>
Date: Fri, 8 Mar 2024 12:57:51 +0100
Subject: [PATCH] 2nd draft of first interaction reader
---
applications/c8_air_shower.cpp | 6 +-
.../modules/writers/InteractionWriter.inl | 71 +++++++++-----
corsika/modules/writers/EnergyLossWriter.hpp | 2 +-
corsika/modules/writers/InteractionWriter.hpp | 9 +-
python/corsika/io/outputs/__init__.py | 1 +
python/corsika/io/outputs/interaction.py | 90 ++++++++++++++++++
python/examples/first_interactions.py | 93 +++++++++++++++++++
tests/output/testInteractionWriter.cpp | 13 ++-
8 files changed, 256 insertions(+), 29 deletions(-)
create mode 100644 python/corsika/io/outputs/interaction.py
create mode 100644 python/examples/first_interactions.py
diff --git a/applications/c8_air_shower.cpp b/applications/c8_air_shower.cpp
index 7abd71bbb..ff61c9a9d 100644
--- a/applications/c8_air_shower.cpp
+++ b/applications/c8_air_shower.cpp
@@ -27,6 +27,7 @@
#include <corsika/framework/utility/SaveBoostHistogram.hpp>
#include <corsika/modules/writers/EnergyLossWriter.hpp>
+#include <corsika/modules/writers/InteractionWriter.hpp>
#include <corsika/modules/writers/LongitudinalWriter.hpp>
#include <corsika/modules/writers/PrimaryWriter.hpp>
#include <corsika/modules/writers/SubWriter.hpp>
@@ -571,10 +572,13 @@ int main(int argc, char** argv) {
// register ZHS with the output manager
output.add("ZHS", zhs);
+ InteractionWriter<setup::Tracking, ParticleWriterParquet> inter_writer(showerAxis, observationLevel);
+ output.add("interactions", inter_writer);
+
// assemble the final process sequence with radio
auto sequence = make_sequence(stackInspect, neutrinoPrimaryPythia, hadronSequence,
decayPythia, emCascade, emContinuous, coreas, zhs,
- longprof, observationLevel, thinning, cut);
+ longprof, observationLevel, inter_writer, thinning, cut);
/* === END: SETUP PROCESS LIST === */
diff --git a/corsika/detail/modules/writers/InteractionWriter.inl b/corsika/detail/modules/writers/InteractionWriter.inl
index 61784ae1c..596844035 100644
--- a/corsika/detail/modules/writers/InteractionWriter.inl
+++ b/corsika/detail/modules/writers/InteractionWriter.inl
@@ -12,15 +12,18 @@
namespace corsika {
- inline InteractionWriter::InteractionWriter(ShowerAxis const& axis)
- : showerAxis_(axis)
+ template <typename TTracking, typename TOutput>
+ inline InteractionWriter<TTracking, TOutput>::InteractionWriter(ShowerAxis const& axis, ObservationPlane<TTracking, TOutput> const& obsPlane)
+ : obsPlane_(obsPlane)
+ , showerAxis_(axis)
, interactionCounter_(0)
, showerId_(0)
, nSecondaries_(0)
{}
+ template <typename TTracking, typename TOutput>
template <typename TStackView>
- inline void InteractionWriter::doSecondaries(TStackView& vS) {
+ inline void InteractionWriter<TTracking, TOutput>::doSecondaries(TStackView& vS) {
if (interactionCounter_++) {
return;
}
@@ -30,12 +33,17 @@ namespace corsika {
CORSIKA_LOG_INFO("First interaction at dX {}", dX);
CORSIKA_LOG_INFO("Primary: {}, E {}", primary.getPID(), primary.getKineticEnergy());
+ Vector const dr_prim = primary.getPosition() - obsPlane_.getPlane().getCenter();
+ auto const p_prim = primary.getMomentum();
+
*(output_.getWriter()) << showerId_ << true
- << static_cast<int>(primary.getPID())
- << static_cast<float>(primary.getKineticEnergy() / 1_GeV)
- << static_cast<float>(1.0)
- << static_cast<float>(0.0)
- << static_cast<float>(0.0)
+ << 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;
@@ -44,12 +52,17 @@ namespace corsika {
auto particle = vS.begin();
while (particle != vS.end()) {
+ Vector const dr_2nd = particle.getPosition() - obsPlane_.getPlane().getCenter();
+ auto const p_2nd = particle.getMomentum();
+
*(output_.getWriter()) << showerId_ << false
- << static_cast<int>(particle.getPID())
- << static_cast<float>(particle.getKineticEnergy() / 1_GeV)
- << static_cast<float>(1.0)
- << static_cast<float>(0.0)
- << static_cast<float>(0.0)
+ << 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)
+ << 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;
@@ -60,20 +73,25 @@ namespace corsika {
}
}
- inline void InteractionWriter::startOfLibrary(boost::filesystem::path const& directory) {
+ template <typename TTracking, typename TOutput>
+ inline void InteractionWriter<TTracking, TOutput>::startOfLibrary(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("kinetic_energy", parquet::Repetition::REQUIRED,
- parquet::Type::FLOAT, parquet::ConvertedType::NONE);
- output_.addField("nx", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
+ 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("ny", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
+ output_.addField("py", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
parquet::ConvertedType::NONE);
- output_.addField("nz", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
+ output_.addField("pz", parquet::Repetition::REQUIRED, parquet::Type::FLOAT,
parquet::ConvertedType::NONE);
output_.addField("time", parquet::Repetition::REQUIRED, parquet::Type::DOUBLE,
parquet::ConvertedType::NONE);
@@ -88,21 +106,25 @@ namespace corsika {
summary_ = YAML::Node();
}
- inline void InteractionWriter::startOfShower(unsigned int const showerId) {
+ template <typename TTracking, typename TOutput>
+ inline void InteractionWriter<TTracking, TOutput>::startOfShower(unsigned int const showerId) {
showerId_ = showerId;
interactionCounter_ = 0;
nSecondaries_ = 0;
}
- inline void InteractionWriter::endOfShower(unsigned int const) {
+ 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::endOfLibrary() {
+ template <typename TTracking, typename TOutput>
+ inline void InteractionWriter<TTracking, TOutput>::endOfLibrary() {
output_.closeStreamer();
}
- inline YAML::Node InteractionWriter::getConfig() const {
+ template <typename TTracking, typename TOutput>
+ inline YAML::Node InteractionWriter<TTracking, TOutput>::getConfig() const {
YAML::Node node;
node["type"] = "Interactions";
node["units"]["energy"] = "GeV";
@@ -112,7 +134,8 @@ namespace corsika {
return node;
}
- inline YAML::Node InteractionWriter::getSummary() const {
+ template <typename TTracking, typename TOutput>
+ inline YAML::Node InteractionWriter<TTracking, TOutput>::getSummary() const {
return summary_;
}
diff --git a/corsika/modules/writers/EnergyLossWriter.hpp b/corsika/modules/writers/EnergyLossWriter.hpp
index 6f42de944..d0a201062 100644
--- a/corsika/modules/writers/EnergyLossWriter.hpp
+++ b/corsika/modules/writers/EnergyLossWriter.hpp
@@ -161,7 +161,7 @@ namespace corsika {
/**
* Return a summary.
*/
- YAML::Node getSummary() const;
+ YAML::Node getSummary() const override;
/**
* Return the configuration of this output.
diff --git a/corsika/modules/writers/InteractionWriter.hpp b/corsika/modules/writers/InteractionWriter.hpp
index 6ba3fee30..6c1509da2 100644
--- a/corsika/modules/writers/InteractionWriter.hpp
+++ b/corsika/modules/writers/InteractionWriter.hpp
@@ -12,19 +12,22 @@
#include <corsika/media/ShowerAxis.hpp>
+#include <corsika/modules/ObservationPlane.hpp>
+
#include <corsika/output/BaseOutput.hpp>
#include <corsika/output/ParquetStreamer.hpp>
namespace corsika {
- class InteractionWriter : public SecondariesProcess<InteractionWriter>, BaseOutput{
+ template <typename TTracking, typename TOutput>
+ class InteractionWriter : public SecondariesProcess<InteractionWriter<TTracking, TOutput>>, public BaseOutput{
public:
/**
*
* @param axis - axis used for calculating grammage
*/
- InteractionWriter(ShowerAxis const& axis);
+ InteractionWriter(ShowerAxis const& axis, ObservationPlane<TTracking, TOutput> const& obsPlane);
/**
*
@@ -46,7 +49,9 @@ namespace corsika {
auto getInteractionCounter() { return interactionCounter_; }
private:
+ ObservationPlane<TTracking, TOutput> const obsPlane_;
ShowerAxis const& showerAxis_;
+
unsigned int interactionCounter_;
unsigned int showerId_;
unsigned int nSecondaries_;
diff --git a/python/corsika/io/outputs/__init__.py b/python/corsika/io/outputs/__init__.py
index 54731dd42..7b35c005b 100644
--- a/python/corsika/io/outputs/__init__.py
+++ b/python/corsika/io/outputs/__init__.py
@@ -16,6 +16,7 @@ from .particle_cut import ParticleCut
from .primary import Particle, PrimaryParticle
from .radio_process import RadioProcess
from .track_writer import TrackWriter
+from .interaction import Interactions
__all__ = [
"Output",
diff --git a/python/corsika/io/outputs/interaction.py b/python/corsika/io/outputs/interaction.py
new file mode 100644
index 000000000..aabe38907
--- /dev/null
+++ b/python/corsika/io/outputs/interaction.py
@@ -0,0 +1,90 @@
+"""
+ Read data written by EnergyLoss.
+
+ (c) Copyright 2024 CORSIKA Project, corsika-project@lists.kit.edu
+
+ This software is distributed under the terms of the GNU General Public
+ Licence version 3 (GPL Version 3). See file LICENSE for a full version of
+ the license.
+"""
+
+import logging
+import os.path as op
+from typing import Any
+
+import pyarrow.parquet as pq
+
+from ..converters import arrow_to_numpy
+from .output import Output
+
+class Interactions(Output):
+ """
+ Reads the interactions and secondaries
+ """
+
+ def __init__(self, path: str):
+ """
+ Load the particle data into a parquet table.
+
+ Parameters
+ ----------
+ path: str
+ The path to the directory containing this output.
+ """
+ super().__init__(path)
+
+ # try and load our data
+ try:
+ self.__data = pq.read_table(op.join(path, "interactions.parquet"))
+ except Exception as e:
+ logging.getLogger("corsika").warn(
+ f"An error occured loading an Interaction: {e}"
+ )
+
+ def is_good(self) -> bool:
+ """
+ Returns true if this output has been read successfully
+ and has the correct files/state/etc.
+
+ Returns
+ -------
+ bool:
+ True if this is a good output.
+ """
+ return self.__data is not None
+
+ def astype(self, dtype: str = "pandas", **kwargs: Any) -> Any:
+ """
+ Load the particle data of the particle interactions.
+
+ All additional keyword arguments are passed to `parquet.read_table`
+
+ Parameters
+ ----------
+ dtype: str
+ The data format to return the data in (i.e. numpy, pandas, etc.)
+
+ Returns
+ -------
+ Any:
+ The return type of this method is determined by `dtype`.
+ """
+ if dtype == "arrow":
+ return self.__data
+ elif dtype == "pandas":
+ return self.__data.to_pandas()
+ elif dtype == "numpy":
+ return arrow_to_numpy.convert_to_numpy(self.__data)
+ else:
+ raise ValueError(
+ (
+ f"Unknown format '{dtype}' for Interaction. "
+ "We currently only support ['arrow', 'pandas', 'numpy']."
+ )
+ )
+
+ def __repr__(self) -> str:
+ """
+ Return a string representation of this class.
+ """
+ return f"Interaction('{self.config['name']}')"
diff --git a/python/examples/first_interactions.py b/python/examples/first_interactions.py
new file mode 100644
index 000000000..ef2635638
--- /dev/null
+++ b/python/examples/first_interactions.py
@@ -0,0 +1,93 @@
+#!/usr/bin/python3
+
+import argparse
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import particle
+
+import corsika
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+ "--input-dir", required=True, help="output directory of the CORSIKA 8 simulation"
+)
+parser.add_argument(
+ "--output-dir",
+ default=os.path.join(here, "example_plots"),
+ help="output directory for plots",
+)
+args = parser.parse_args()
+
+if not os.path.isdir(args.output_dir):
+ print("Making directory", args.output_dir)
+ os.makedirs(args.output_dir)
+
+# Load the shower simulation output
+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
+
+shower_ids = np.unique(interactions["shower"])
+
+ncols = 1
+nrows = len(shower_ids)
+fig, ax = plt.subplots(
+ ncols=ncols,
+ nrows=nrows,
+ figsize=(ncols * 8, nrows * 3),
+ gridspec_kw={"wspace": 0.2, "hspace": 0.3},
+)
+ax = np.atleast_1d(ax).flatten()
+
+for sh_id in shower_ids:
+ data = interactions[interactions["shower"] == sh_id]
+
+ 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]
+
+ 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)
+
+
+ # Get daughters, filter out zero-momentum particles
+ daughters = data[data["primary"] == False]
+ p_tot = np.sqrt(daughters["px"]**2 + daughters["py"]**2 + daughters["pz"]**2)
+ daughters = daughters[p_tot > 0]
+
+ print(daughters)
+
+ for i in range(len(daughters)):
+ pz = daughters["pz"].iloc[i]
+ px = np.cos(-angle) * daughters["px"].iloc[i] + np.sin(-angle) * daughters["py"].iloc[i]
+
+ pdg = int(daughters["pdg"].iloc[i])
+ try:
+ name = f"${particle.Particle.from_pdgid(pdg).latex_name}$"
+ except Exception:
+ name = "Unknown"
+ print(pdg)
+
+ ax[sh_id].plot([px, 0], [pz, 0], label=name)
+
+ ax[sh_id].set_aspect("equal")
+ ax[sh_id].legend()
+
+ 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)
+
+plot_path = os.path.join(args.output_dir, "interactions.png")
+print("Saving", plot_path)
+fig.savefig(plot_path, bbox_inches="tight")
diff --git a/tests/output/testInteractionWriter.cpp b/tests/output/testInteractionWriter.cpp
index 302c0ccd6..8f982a4ba 100644
--- a/tests/output/testInteractionWriter.cpp
+++ b/tests/output/testInteractionWriter.cpp
@@ -20,6 +20,8 @@
#include <corsika/media/Environment.hpp>
#include <corsika/media/HomogeneousMedium.hpp>
+#include <corsika/setup/SetupTrajectory.hpp>
+
#include <SetupTestStack.hpp>
#include <SetupTestTrajectory.hpp>
@@ -61,7 +63,16 @@ TEST_CASE("InteractionWriter", "process") {
ShowerAxis axis(start, length, *env);
- InteractionWriter writer(axis);
+ Plane const plane(Point(rootCS, {0_m, 0_m, 0_m}),
+ DirectionVector(rootCS, {0., 0., 1.}));
+
+ ObservationPlane<setup::Tracking, ParticleWriterParquet> obsPlane{
+ plane, DirectionVector(rootCS, {1., 0., 0.}),
+ true, // plane should "absorb" particles
+ false};
+
+
+ InteractionWriter<setup::Tracking, ParticleWriterParquet> writer(axis, obsPlane);
writer.startOfLibrary(outputDir);
writer.startOfShower(0);
--
GitLab