Merge pull request #21 from cta-observatory/noparse

Add option to not parse blocks

corsikaio/file.py

@@ -24,12 +24,19 @@
     ['header', 'particles', 'longitudinal', 'end']
 )
 
+def _to_floatarray(block):
+    return np.frombuffer(block, dtype=np.float32)
+
 
 class CorsikaFile:
+    """
+    A file to iterate over events in a CORSIKA binary output file
+    """
 
-    def __init__(self, path):
+    def __init__(self, path, parse_blocks=True):
         self.EventClass = Event
 
+        self.parse_blocks = parse_blocks
         self._buffer_size = read_buffer_size(path)
         self._f = open_compressed(path)
 
@@ -72,7 +79,10 @@ def __next__(self):
         if block[:4] != b'EVTH':
             raise IOError('EVTH block expected but found {}'.format(block[:4]))
 
-        event_header = parse_event_header(block)[0]
+        if self.parse_blocks:
+            event_header = parse_event_header(block)[0]
+        else:
+            event_header = _to_floatarray(block)
 
         data_bytes = bytearray()
         long_bytes = bytearray()
@@ -87,9 +97,14 @@ def __next__(self):
 
             block = self.read_block()
 
-        event_end = parse_event_end(block)[0]
-        data = self.parse_data_blocks(data_bytes)
-        longitudinal = parse_longitudinal(long_bytes)
+        if self.parse_blocks:
+            event_end = parse_event_end(block)[0]
+            data = self.parse_data_blocks(data_bytes)
+            longitudinal = parse_longitudinal(long_bytes)
+        else:
+            event_end = _to_floatarray(block)
+            data = _to_floatarray(data_bytes).reshape(-1, 7)
+            longitudinal = _to_floatarray(long_bytes)
 
         return self.EventClass(event_header, data, longitudinal, event_end)
 
@@ -150,8 +165,8 @@ def close(self):
 
 class CorsikaCherenkovFile(CorsikaFile):
 
-    def __init__(self, path, mmcs=False):
-        super().__init__(path)
+    def __init__(self, path, mmcs=False, **kwargs):
+        super().__init__(path, **kwargs)
 
         self.EventClass = PhotonEvent
         self.mmcs = mmcs
@@ -175,8 +190,8 @@ def parse_data_blocks(self, data_bytes):
 
 class CorsikaParticleFile(CorsikaFile):
 
-    def __init__(self, path):
-        super().__init__(path)
+    def __init__(self, path, **kwargs):
+        super().__init__(path, **kwargs)
         self.EventClass = ParticleEvent
 
     def parse_data_blocks(self, data_bytes):

tests/test_file.py

@@ -1,4 +1,5 @@
 import pytest
+import numpy as np
 
 
 def test_version():
@@ -71,6 +72,22 @@ def test_particle_longi():
         assert e.header['event_number'] == 1
 
 
+def test_particle_no_parse():
+    from corsikaio import CorsikaParticleFile
+
+    with CorsikaParticleFile('tests/resources/corsika757_particle', parse_blocks=False) as f:
+        n_read = 0
+        for e in f:
+            n_read += 1
+            # second entry in header is event_number
+            assert e.header[1] == n_read
+            assert e.header.dtype == np.float32
+            assert len(e.header) == 273
+            assert e.particles.size % 273 == 0
+        assert n_read == 10
+
+
+
 def test_truncated(tmp_path):
     '''Test we raise a meaningful error for a truncated file