Merge pull request #23 from cta-observatory/feature_add-field-unit

Add a unit to all fields

corsikaio/subblocks/__init__.py

@@ -13,11 +13,11 @@
 from ..constants import RUNH_VERSION_POSITION, EVTH_VERSION_POSITION
 
 __all__ = [
-    'parse_event_header',
-    'parse_run_header'
-    'parse_cherenkov_photons',
-    'parse_particle_data',
-    'parse_longitudinal',
+    "parse_event_header",
+    "parse_run_header",
+    "parse_cherenkov_photons",
+    "parse_particle_data",
+    "parse_longitudinal",
 ]
 
 
@@ -44,7 +44,7 @@ def parse_event_end(event_end_bytes):
 
 def get_version(header_bytes, version_pos):
     sl = slice(4 * (version_pos - 1), 4 * version_pos)
-    return round(struct.unpack('f', header_bytes[sl])[0], 4)
+    return round(struct.unpack("f", header_bytes[sl])[0], 4)
 
 
 def parse_data_block(data_block_bytes, dtype):
@@ -63,3 +63,23 @@ def parse_particle_data(data_block_bytes):
 
 def parse_longitudinal(longitudinal_data_bytes):
     return parse_data_block(longitudinal_data_bytes, longitudinal_data_dtype)
+
+
+def get_units_from_fields(subblock_fields):
+    """Retrieve units as a dictionary from the fields of a file subblock.
+
+    Dimensionless fields are not selected.
+
+    Parameters
+    ----------
+    subblock_fields: list(Field)
+        One of the defined lists of fields.
+
+    Return
+    ------
+    units: dict
+        Dictionary with field names as keys and
+        string representations of units as values.
+    """
+    units = {field.name: field.unit for field in subblock_fields if field.unit is not None}
+    return units

corsikaio/subblocks/data.py

@@ -1,33 +1,39 @@
-import numpy as np
+from .dtypes import build_dtype, Field
 
-cherenkov_photons_dtype = np.dtype([
-    ('n_photons', 'float32'),
-    ('x', 'float32'),
-    ('y', 'float32'),
-    ('u', 'float32'),
-    ('v', 'float32'),
-    ('t', 'float32'),
-    ('production_height', 'float32'),
-])
+cherenkov_photons_fields = [
+    Field(1, "n_photons", dtype="float32"),
+    Field(2, "x", unit="cm", dtype="float32"),
+    Field(3, "y", unit="cm", dtype="float32"),
+    Field(4, "u", dtype="float32"),
+    Field(5, "v", dtype="float32"),
+    Field(6, "t", unit="ns", dtype="float32"),
+    Field(7, "production_height", unit="cm", dtype="float32"),
+]
 
-mmcs_cherenkov_photons_dtype = np.dtype([
-    ('n_photons', 'float32'),
-    ('x', 'float32'),
-    ('y', 'float32'),
-    ('u', 'float32'),
-    ('v', 'float32'),
-    ('t', 'float32'),
-    ('wavelength', 'float32'),
-    ('mother_particle', 'uint16'),
-    ('production_height', 'float32'),
-])
+cherenkov_photons_dtype = build_dtype(cherenkov_photons_fields, itemsize=None)
 
-particle_data_dtype = np.dtype([
-    ('particle_description', 'float32'),
-    ('px', 'float32'),
-    ('py', 'float32'),
-    ('pz', 'float32'),
-    ('x', 'float32'),
-    ('y', 'float32'),
-    ('t', 'float32'),
-])
+mmcs_cherenkov_photons_fields = [
+    Field(1, "n_photons", dtype="float32"),
+    Field(2, "x", unit="cm", dtype="float32"),
+    Field(3, "y", unit="cm", dtype="float32"),
+    Field(4, "u", dtype="float32"),
+    Field(5, "v", dtype="float32"),
+    Field(6, "t", unit="ns", dtype="float32"),
+    Field(7, "wavelength", unit="nm", dtype="float32"),
+    Field(8, "mother_particle", dtype="uint16"),
+    Field(9, "production_height", unit="cm", dtype="float32"),
+]
+
+mmcs_cherenkov_photons_dtype = build_dtype(mmcs_cherenkov_photons_fields, itemsize=None)
+
+particle_data_fields = [
+    Field(1, "particle_description", dtype="float32"),
+    Field(2, "px", unit="GeV/c", dtype="float32"),
+    Field(3, "py", unit="GeV/c", dtype="float32"),
+    Field(4, "pz", unit="GeV/c", dtype="float32"),
+    Field(5, "x", unit="cm", dtype="float32"),
+    Field(6, "y", unit="cm", dtype="float32"),
+    Field(7, "t", unit="ns", dtype="float32"),
+]
+
+particle_data_dtype = build_dtype(particle_data_fields, itemsize=None)

corsikaio/subblocks/dtypes.py

@@ -2,22 +2,24 @@
 import numpy as np
 
 
-Field = namedtuple('Field', ['position', 'name', 'shape', 'dtype'])
+Field = namedtuple("Field", ["position", "name", "unit", "shape", "dtype"])
 
-# set defaults for shape and dtype, this works for the 2 right most attributes
-Field.__new__.__defaults__ = (1, 'f4')
+# set defaults for unit, shape and dtype, this works for the 3 right most attributes
+Field.__new__.__defaults__ = (None, 1, "f4")
 
 
 def build_dtype(fields, itemsize=4 * 273):
     dt = defaultdict(list)
-    dt['itemsize'] = itemsize
+    if itemsize is not None:
+        dt["itemsize"] = itemsize
 
     for field in fields:
-        dt['names'].append(field.name)
-        dt['offsets'].append((field.position - 1) * 4)
+        dt["names"].append(field.name)
+        dt["offsets"].append((field.position - 1) * 4)
         if field.shape != 1:
-            dt['formats'].append((field.dtype, field.shape))
+            dt["formats"].append((field.dtype, field.shape))
         else:
-            dt['formats'].append(field.dtype)
+            dt["formats"].append(field.dtype)
 
     return np.dtype(dict(**dt))
+

corsikaio/subblocks/event_end.py

@@ -1,18 +1,20 @@
 from .dtypes import build_dtype, Field
 
 
-event_end_dtype = build_dtype([
-    Field(1, 'event_end', dtype='S4'),
-    Field(2, 'event_number'),
-    Field(3, 'n_photons_weighted'),
-    Field(4, 'n_electrons_weighted'),
-    Field(5, 'n_hadrons_weighted'),
-    Field(6, 'n_muons_weighted'),
-    Field(7, 'n_particles_written'),
-    Field(262, 'chi_square_longitudinal'),
-    Field(263, 'n_photons_written'),
-    Field(264, 'n_electrons_written'),
-    Field(265, 'n_hadrons_written'),
-    Field(266, 'n_muons_written'),
-    Field(267, 'n_em_particles_preshower'),
-])
+event_end_fields = [
+    Field(1, "event_end", dtype="S4"),
+    Field(2, "event_number"),
+    Field(3, "n_photons_weighted"),
+    Field(4, "n_electrons_weighted"),
+    Field(5, "n_hadrons_weighted"),
+    Field(6, "n_muons_weighted"),
+    Field(7, "n_particles_written"),
+    Field(262, "chi_square_longitudinal"),
+    Field(263, "n_photons_written"),
+    Field(264, "n_electrons_written"),
+    Field(265, "n_hadrons_written"),
+    Field(266, "n_muons_written"),
+    Field(267, "n_em_particles_preshower"),
+]
+
+event_end_dtype = build_dtype(event_end_fields)