diff --git a/python/corsika/io/outputs/__init__.py b/python/corsika/io/outputs/__init__.py index d39d9a8b4bd68465ab1d799f10c15bb20032cbc9..63b76b533d41085aa77af57df5439c5dd3c68335 100644 --- a/python/corsika/io/outputs/__init__.py +++ b/python/corsika/io/outputs/__init__.py @@ -12,6 +12,7 @@ from .track_writer import TrackWriter from .longitudinal_profile import LongitudinalProfile from .bethe_bloch import BetheBlochPDG from .particle_cut import ParticleCut +from .energy_loss import EnergyLoss from .output import Output __all__ = [ @@ -21,4 +22,5 @@ __all__ = [ "LongitudinalProfile", "BetheBlochPDG", "ParticleCut", + "EnergyLoss" ] diff --git a/python/corsika/io/outputs/energy_loss.py b/python/corsika/io/outputs/energy_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..76846dd2c83d497f68aaf38d3f3e4823b91ee8a3 --- /dev/null +++ b/python/corsika/io/outputs/energy_loss.py @@ -0,0 +1,87 @@ +""" + Read data written by EnergyLoss. + + (c) Copyright 2020 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 .output import Output + + +class EnergyLoss(Output): + """ + Read particle data from a EnergyLoss process. + """ + + 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, "dEdX.parquet")) + except Exception as e: + logging.getLogger("corsika").warn( + f"An error occured loading a EnergyLoss: {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 from this track writer. + + 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() + else: + raise ValueError( + ( + f"Unknown format '{dtype}' for EnergyLoss. " + "We currently only support ['arrow', 'pandas']." + ) + ) + + def __repr__(self) -> str: + """ + Return a string representation of this class. + """ + return f"EnergyLess('{self.config['name']}')"