From c2cf8258ccdfa0a88d32dc352f167a170a8b6cee Mon Sep 17 00:00:00 2001
From: Alan Coleman <alanc@udel.edu>
Date: Wed, 13 Mar 2024 22:07:51 +0100
Subject: [PATCH] update python style
---
python/corsika/io/outputs/interaction.py | 11 ++++++++---
python/examples/first_interactions.py | 6 ++++--
2 files changed, 12 insertions(+), 5 deletions(-)
diff --git a/python/corsika/io/outputs/interaction.py b/python/corsika/io/outputs/interaction.py
index e721243a8..21a1b1d54 100644
--- a/python/corsika/io/outputs/interaction.py
+++ b/python/corsika/io/outputs/interaction.py
@@ -9,22 +9,24 @@
"""
import logging
-import pyarrow.parquet as pq
import os.path as op
from typing import Any
-import yaml
+import pyarrow.parquet as pq
+import yaml
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
@@ -34,6 +36,9 @@ class FirstInteraction(Particle):
"""
Particle.__init__(self, prop_dict)
+ # need to define for momentum prop
+ self.px = self.py = self.pz = None
+
def __repr__(self) -> str:
"""
Return a string representation of this class.
@@ -43,11 +48,11 @@ class FirstInteraction(Particle):
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
diff --git a/python/examples/first_interactions.py b/python/examples/first_interactions.py
index de433f1a5..d86d8d70d 100644
--- a/python/examples/first_interactions.py
+++ b/python/examples/first_interactions.py
@@ -54,7 +54,7 @@ for ish, sh_id in enumerate(shower_ids):
daughters = interactions[interactions["shower"] == sh_id]
mother = projectiles[ish]
- norm = np.sqrt(sum(np.array(mother.momentum)**2)) # total momentum
+ norm = np.sqrt(sum(np.array(mother.momentum) ** 2)) # total momentum
print("Plotting shower", sh_id)
print(mother)
@@ -69,9 +69,11 @@ for ish, sh_id in enumerate(shower_ids):
p_tot = np.sqrt(daughters["px"] ** 2 + daughters["py"] ** 2 + daughters["pz"] ** 2)
daughters = daughters[p_tot > 0]
+ # keep track of largest values
max_pz = abs(pz)
max_pxy = max(abs(px), abs(py))
+ # plot the daughter particles
for i in range(len(daughters)):
pz = daughters["pz"].iloc[i]
px = daughters["px"].iloc[i]
@@ -89,7 +91,7 @@ for ish, sh_id in enumerate(shower_ids):
max_pz = max(max_pz, abs(pz))
max_pxy = max(max_pxy, max(abs(px), abs(py)))
-
+ # set plotting style and names
for i in range(2):
ax[i, sh_id].legend()
--
GitLab