722 Integrated buffer writer

#722

[author: andLaing]

Adds a new io option to `rwf_io` which uses the basic `rwf_writer` and
other table writers to write all event info in one step. In this way
long MC events can be split into multiple trigger-like events in the
output file and the event numbers can be logged and mapped
accordingly.

Some issues remain for the logging portion that are under debate.

Addresses point 2 of issue #691

[reviewer: jmalbos]

This PR adds a new writer (and the corresponding test) to `rwf_io` than
can handle the splitting of long (MC) events into several subevents. A
new table (`MCEventMap`) is used to associate the new subevents to the
original event.

Nevertheless, this new writer has limited use until a decision is
taken regarding #715.

invisible_cities/evm/nh5.py

@@ -43,6 +43,12 @@ class MCExtentInfo(tb.IsDescription):
     last_particle = tb.UInt64Col(pos=2)
 
 
+class MCEventMap(tb.IsDescription):
+    """Map between event index and original event."""
+    evt_number = tb.Int32Col(shape=(), pos=0)
+    sub_evt    = tb.Int32Col(shape=(), pos=1)
+
+
 class MCHitInfo(tb.IsDescription):
     """Stores the simulated hits as metadata using Pytables.
     """

invisible_cities/io/rwf_io.py

@@ -1,9 +1,16 @@
 import numpy  as np
 import tables as tb
 
-from typing  import Callable
+from functools import  partial
+from typing    import Callable
+from typing    import     List
+from typing    import Optional
+from typing    import    Tuple
 
-from .. reco import tbl_functions as tbl
+from .. evm .nh5         import           MCEventMap
+from .. reco             import        tbl_functions as tbl
+from .  run_and_event_io import run_and_event_writer
+from .  table_io         import           make_table
 
 
 def rwf_writer(h5out           : tb.file.File          ,
@@ -54,3 +61,97 @@ def write_rwf(waveform : np.ndarray) -> None:
         """
         rwf_table.append(waveform.reshape(1, n_sensors, waveform_length))
     return write_rwf
+
+
+def buffer_writer(h5out, *,
+                  run_number :          int           ,
+                  n_sens_eng :          int           ,
+                  n_sens_trk :          int           ,
+                  length_eng :          int           ,
+                  length_trk :          int           ,
+                  group_name : Optional[str] =    None,
+                  compression: Optional[str] = 'ZLIB4'
+                  ) -> Callable[[int, List, List], None]:
+    """
+    Generalised buffer writer which defines a raw waveform writer
+    for each type of sensor as well as an event info writer.
+    Each call gives a list of 'triggers' to be written as
+    separate events in the output.
+
+    parameters
+    ----------
+    run_number  : int
+                  Run number to be saved in runInfo.
+    n_sens_eng  : int
+                  Number of sensors in the energy plane.
+    n_sens_trk  : int
+                  Number of sensors in the tracking plane.
+    length_eng  : int
+                  Number of samples per waveform for energy plane.
+    length_trk  : int
+                  Number of samples per waveform for tracking plane.
+    group_name  : Optional[str] default None
+                  Group name within root where waveforms to be saved.
+                  Default directly in root
+    compression : Optional[str] default 'ZLIB4'
+                  Compression level for output file.
+
+    returns
+    -------
+    write_buffers : Callable
+                    A function which takes event information
+                    for the tracking and energy planes and
+                    the event timestamps and saves to file.
+    """
+
+    eng_writer = rwf_writer(h5out,
+                            group_name      =  group_name,
+                            compression     = compression,
+                            table_name      =     'pmtrd',
+                            n_sensors       =  n_sens_eng,
+                            waveform_length =  length_eng)
+
+    trk_writer = rwf_writer(h5out,
+                            group_name      =  group_name,
+                            compression     = compression,
+                            table_name      =    'sipmrd',
+                            n_sensors       =  n_sens_trk,
+                            waveform_length =  length_trk)
+
+    run_and_event = partial(run_and_event_writer(h5out                  ,
+                                                 compression=compression),
+                            run_number = run_number                      )
+
+    nexus_map = make_table(h5out, 'Run', 'eventMap', MCEventMap,
+                           "event & nexus evt for each index", compression)
+
+    def write_buffers(nexus_evt  :        int ,
+                      timestamps : List[  int],
+                      events     : List[Tuple]) -> None:
+        """
+        Write run info and event waveforms to file.
+        
+        parameters
+        ----------
+        nexus_evt  :  int
+                     Event number from MC output file.
+        timestamps : List[int]
+                     List of event times
+        events     : List[Tuple]
+                     List of tuples containing the energy and
+                     tracking plane info for each identified 'trigger'.
+        """
+
+        for i, (t_stamp, (eng, trk)) in enumerate(zip(timestamps, events)):
+            ## The exact way to log MC event splitting
+            ## still to be decided.
+            run_and_event(event_number=nexus_evt, timestamp=t_stamp)
+            mrow = nexus_map.row
+            mrow["evt_number"] = nexus_evt
+            mrow[   "sub_evt"] = i
+            mrow.append()
+            ##
+
+            eng_writer(eng)
+            trk_writer(trk)
+    return write_buffers

invisible_cities/io/rwf_io_test.py

@@ -3,9 +3,11 @@
 import numpy  as np
 import tables as tb
 
-from pytest   import       mark
+from pytest   import       fixture
+from pytest   import          mark
 
-from . rwf_io import rwf_writer
+from . rwf_io import    rwf_writer
+from . rwf_io import buffer_writer
 
 
 @mark.parametrize("group_name", (None, 'RD', 'BLR'))
@@ -42,3 +44,51 @@ def test_rwf_writer(config_tmpdir, group_name):
         table = getattr(group, table_name)
         assert table.shape == (nevt, nsensor, nsample)
         assert np.all(test_data == table.read())
+
+
+@mark.parametrize("event triggers".split(),
+                  ((2, [10]), (3, [10, 1100]), (4, [20])))
+def test_buffer_writer(config_tmpdir, event, triggers):
+
+    run_number  = -6400
+    n_pmt       =    12
+    nsamp_pmt   =   100
+    n_sipm      =  1792
+    nsamp_sipm  =    10
+
+    buffers = [(np.random.poisson(5, (n_pmt , nsamp_pmt )),
+                np.random.poisson(5, (n_sipm, nsamp_sipm))) for _ in triggers]
+
+    out_name = os.path.join(config_tmpdir, 'test_buffers.h5')
+    with tb.open_file(out_name, 'w') as data_out:
+
+        buffer_writer_ = buffer_writer(data_out,
+                                       run_number = run_number,
+                                       n_sens_eng =      n_pmt,
+                                       n_sens_trk =     n_sipm,
+                                       length_eng =  nsamp_pmt,
+                                       length_trk = nsamp_sipm)
+
+        buffer_writer_(event, triggers, buffers)
+
+    pmt_wf  = np.array([b[0] for b in buffers])
+    sipm_wf = np.array([b[1] for b in buffers])
+    with tb.open_file(out_name) as h5saved:
+        assert 'Run'      in h5saved.root
+        assert 'pmtrd'    in h5saved.root
+        assert 'sipmrd'   in h5saved.root
+        assert 'events'   in h5saved.root.Run
+        assert 'runInfo'  in h5saved.root.Run
+        assert 'eventMap' in h5saved.root.Run
+
+        nsaves = len(triggers)
+        assert len(h5saved.root.Run.events  ) == nsaves
+        assert len(h5saved.root.Run.eventMap) == nsaves
+        assert len(h5saved.root.Run.runInfo ) == nsaves
+        assert np.all([r[0] == run_number for r in h5saved.root.Run.runInfo])
+
+        assert h5saved.root.pmtrd .shape == (nsaves, n_pmt ,  nsamp_pmt)
+        assert np.all(h5saved.root.pmtrd [:] ==  pmt_wf)
+
+        assert h5saved.root.sipmrd.shape == (nsaves, n_sipm, nsamp_sipm)
+        assert np.all(h5saved.root.sipmrd[:] == sipm_wf)