Loopplugin touching windows + plugin documentation

docs/source/advanced/plugin_dev.rst

@@ -6,7 +6,7 @@ Special time fields
 The ``time``, ``endtime``, ``dt`` and ``length`` fields have special meaning for strax.
 
 It is useful for most plugins to output a ``time`` and ``endtime`` field, indicating the
-start and (exclusive) end time of the entitities you are producing.
+start and (exclusive) end time of the entities you are producing.
 If you do not do this, your plugin cannot be loaded for part of a run (e.g. with ``seconds_range``).
 
 Both ``time`` and ``endtime`` should be 64-bit integer timestamps in nanoseconds since the unix epoch. Instead of ``endtime``, you can provide ``dt`` (an integer time resolution in ns) and ``length`` (integer); strax will then compute the endtime as ``dt * length``. Lower-level datatypes commonly use this.
@@ -28,7 +28,7 @@ To return multiple outputs from a plugin:
 Options and defaults
 ----------------------
 
-You can specify options using the `strax.takes_config` decorator and the `strax.Option` objects. See any plugin source code for example (todo: don't be lazy and explain).
+You can specify options using the ``strax.takes_config`` decorator and the ``strax.Option`` objects. See any plugin source code for example (todo: don't be lazy and explain).
 
 There is a single configuration dictionary in a strax context, shared by all plugins. Be judicious in how you name your options to avoid clashes. "Threshold" is probably a bad name, "peak_min_channels" is better.
 
@@ -40,25 +40,158 @@ You can specify defaults in several ways:
 
 - ``default``: Use the given value as default.
 - ``default_factory``: Call the given function (with no arguments) to produce a default. Use for mutable values such as lists.
-- ``default_per_run``: Specify a list of 2-tuples: ``(start_run, default)``. Here start_run is a numerized run name (e.g 170118_1327; note the underscore is valid in integers since python 3.6) and ``default`` the option that applies from that run onwards.
+- ``default_per_run``: Specify a list of 2-tuples: ``(start_run, default)``. Here start_run is a numerized run name (e.g ``170118_1327``; note the underscore is valid in integers since python 3.6) and ``default`` the option that applies from that run onwards.
 - The ``strax_defaults`` dictionary in the run metadata. This overrides any defaults specified in the plugin code, but take care -- if you change a value here, there will be no record anywhere of what value was used previously, so you cannot reproduce your results anymore!
 
 
 Plugin types
 ----------------------
 
 There are several plugin types:
-   * `Plugin`: The general type of plugin. Should contain at least `depends_on = <datakind>`, `provides = <datatype>`, `def compute(self, <datakind>)`, and `dtype = <dtype> ` or `def infer_dtype(): <>`.
-   * `OverlapWindowPlugin`: Allows a plugin to look for data in adjacent chunks. A OverlapWindowPlugin assumes: all inputs are sorted by *endtime*. This only works for disjoint intervals such as peaks or events, but NOT records! The user has to define get_window_size(self) along with the plugin which returns the required chunk extension in nanoseconds. 
-   * `LoopPlugin`: Allows user to loop over a given datakind and find the corresponding data of a lower datakind using for example `def compute_loop(self, events, peaks)` where we loop over events and get the corresponding peaks that are within the time range of the event. Currently the second argument (`peaks`) must be fully contained in the first argument (`events` ).
-   * `CutPlugin`: Plugin type where using `def cut_by(self, <datakind>)` inside the plugin a user can return a boolean array that can be used to select data.
-   * `MergeOnlyPlugin`: This is for internal use and only merges two plugins into a new one. See as an example in straxen the `EventInfo` plugin where the following datatypes are merged `'events', 'event_basics', 'event_positions', 'corrected_areas', 'energy_estimates'`.
-   * `ParallelSourcePlugin`: For internal use only to parallelize the processing of low level plugins. This can be activated using stating `parallel = 'process'` in a plugin.
+   * ``Plugin``: The general type of plugin. Should contain at least ``depends_on = <datakind>``, ``provides = <datatype>``, ``def compute(self, <datakind>)``, and ``dtype = <dtype>`` or ``def infer_dtype(): <>``.
+   * ``OverlapWindowPlugin``: Allows a plugin to look for data in adjacent chunks. A ``OverlapWindowPlugin`` assumes all inputs are sorted by *endtime*. This only works for disjoint intervals such as peaks or events, but NOT records! The user has to define ``get_window_size(self)`` along with the plugin which returns the required chunk extension in nanoseconds.
+   * ``LoopPlugin``: Allows user to loop over a given datakind and find the corresponding data of a lower datakind using for example `def compute_loop(self, events, peaks)` where we loop over events and get the corresponding peaks that are within the time range of the event. By default the second argument (``peaks``) must be fully contained in the first argument (``events`` ). If a touching time window is desired set the class attribute ``time_selection`` to `'`touching'``.
+   * ``CutPlugin``: Plugin type where using ``def cut_by(self, <datakind>)`` inside the plugin a user can return a boolean array that can be used to select data.
+   * ``MergeOnlyPlugin``: This is for internal use and only merges two plugins into a new one. See as an example in straxen the ``EventInfo`` plugin where the following datatypes are merged ``'events', 'event_basics', 'event_positions', 'corrected_areas', 'energy_estimates'``.
+   * ``ParallelSourcePlugin``: For internal use only to parallelize the processing of low level plugins. This can be activated using stating ``parallel = 'process'`` in a plugin.
+
+
+Minimal examples
+----------------------
+Below, each of the plugins is minimally worked out, each plugin can be worked
+out into much greater detail, see e.g. the
+`plugins in straxen <https://github.com/XENONnT/straxen/tree/master/straxen/plugins>`_.
+
+strax.Plugin
+____________
+.. code-block:: python
+
+    # To tests, one can use these dummy Peaks and Records from strax
+    import strax
+    import numpy as np
+    from strax.testutils import Records, Peaks, run_id
+    st = strax.Context(register=[Records, Peaks])
+
+    class BasePlugin(strax.Plugin):
+        """The most common plugin where computations on data are performed in strax"""
+        depends_on = 'records'
+
+        # For good practice always specify the version and provide argument
+        provides = 'simple_data'
+        __version__ = '0.0.0'
+
+        # We need to specify the datatype, for this example, we are
+        # going to calculate some areas
+        dtype = strax.time_fields + [(("Total ADC counts",'area'), np.int32)]
+
+        def compute(self, records):
+            result = np.zeros(len(records), dtype=self.dtype)
+
+            # All data in strax must have some sort of time fields
+            result['time'] = records['time']
+            result['endtime'] = strax.endtime(records)
+
+            # For this example, we calculate the total sum of the records-data
+            result['area'] = np.sum(records['data'], axis = 1)
+            return result
+
+    st.register(BasePlugin)
+    st.get_df(run_id, 'simple_data')
+
+
+strax.OverlapWindowPlugin
+_________________________
+.. code-block:: python
+
+    class OverlapPlugin(strax.OverlapWindowPlugin):
+        """
+        Allow peaks get_window_size() left and right to get peaks
+            within the time range
+        """
+        depends_on = 'peaks'
+        provides = 'overlap_data'
+
+        dtype = strax.time_fields + [(("total peaks", 'n_peaks'), np.int16)]
+
+        def get_window_size(self):
+            # Look 10 ns left and right of each peak
+            return 10
+
+        def compute(self, peaks):
+            result = np.zeros(1, dtype=self.dtype)
+            result['time'] = np.min(peaks['time'])
+            result['endtime'] = np.max(strax.endtime(peaks))
+            result['n_peaks'] = len(peaks)
+            return result
+
+    st.register(OverlapPlugin)
+    st.get_df(run_id, 'overlap_data')
+
+
+strax.LoopPlugin
+__________
+.. code-block:: python
+
+    class LoopData(strax.LoopPlugin):
+        """Loop over peaks and find the records within each of those peaks."""
+        depends_on = 'peaks', 'records'
+        provides = 'looped_data'
+
+        dtype = strax.time_fields + [(("total records", 'n_records'), np.int16)]
+
+        # The LoopPlugin specific requirements
+        time_selection = 'fully_contained' # other option is 'touching'
+        loop_over = 'peaks'
+
+        # Use the compute_loop() instead of compute()
+        def compute_loop(self, peaks, records):
+            result = np.zeros(len(peaks), dtype=self.dtype)
+            result['time'] = np.min(peaks['time'])
+            result['endtime'] = np.max(strax.endtime(peaks))
+            result['n_records'] = len(records)
+            return result
+    st.register(LoopData)
+    st.get_df(run_id, 'looped_data')
+
+
+strax.CutPlugin
+_________________________
+.. code-block:: python
+
+    class CutData(strax.CutPlugin):
+        """
+        Create a boolean array if an entry passes a given cut,
+            in this case if the peak has a positive area
+        """
+        depends_on = 'peaks'
+        provides = 'cut_data'
+
+        # Use cut_by() instead of compute() to generate a boolean array
+        def cut_by(self, peaks):
+            return peaks['area']>0
+
+    st.register(CutData)
+    st.get_df(run_id, 'cut_data')
+
+
+strax.MergeOnlyPlugin
+________
+.. code-block:: python
+
+    class MergeData(strax.MergeOnlyPlugin):
+        """Merge datatypes of the same datakind into a single datatype"""
+        depends_on = ('peaks', 'cut_data')
+        provides = 'merged_data'
+
+        # You only need specify the dependencies, those are merged.
+
+    st.register(MergeData)
+    st.get_array(run_id, 'merged_data')
 
 
 Plugin inheritance
 ----------------------
-It is possible to inherit the `compute()` method of an already existing plugin with another plugin. We call these types of plugins child plugins. Child plugins are recognized by strax when the `child_plugin` attribute of the plugin is set to `True`. Below you can find a simple example of a child plugin with its parent plugin:
+It is possible to inherit the ``compute()`` method of an already existing plugin with another plugin. We call these types of plugins child plugins. Child plugins are recognized by strax when the ``child_plugin`` attribute of the plugin is set to ``True``. Below you can find a simple example of a child plugin with its parent plugin:
 
 .. code-block:: python
 
@@ -103,10 +236,10 @@ It is possible to inherit the `compute()` method of an already existing plugin w
             res['width'] = self.config['option_unique_child']
             return res
 
-The `super().compute()` statement in the `compute` method of `ChildPlugin` allows us to execute the code of the parent's compute method without duplicating it. Additionally, if needed, we can extend the code with some for the child-plugin unique computation steps.
+The ``super().compute()`` statement in the ``compute`` method of ``ChildPlugin`` allows us to execute the code of the parent's compute method without duplicating it. Additionally, if needed, we can extend the code with some for the child-plugin unique computation steps.
 
-To allow for the child plugin to have different settings then its parent (e.g. `'by_child_overwrite_option'` in `self.config['by_child_overwrite_option']` of the parent's `compute` method), we have to use specific child option. These options will be recognized by strax and overwrite the config values of the parent parameter during the initialization of the child-plugin. Hence, these changes only affect the child, but not the parent.
+To allow for the child plugin to have different settings then its parent (e.g. ``'by_child_overwrite_option'`` in ``self.config['by_child_overwrite_option']`` of the parent's ``compute`` method), we have to use specific child option. These options will be recognized by strax and overwrite the config values of the parent parameter during the initialization of the child-plugin. Hence, these changes only affect the child, but not the parent.
 
-An option can be flagged as a child option if the corresponding option attribute is set `child_option=True`. Further, the option name which should be overwritten must be specified via the option attribute `parent_option_name`.
+An option can be flagged as a child option if the corresponding option attribute is set ``child_option=True``. Further, the option name which should be overwritten must be specified via the option attribute ``parent_option_name``.
 
 The lineage of a child plugin contains in addition to its options the name and version of the parent plugin.

strax/context.py

@@ -1050,7 +1050,7 @@ def apply_selection(self, x,
         :param x: Numpy structured array
         :param selection_str: Query string or sequence of strings to apply.
         :param time_range: (start, stop) range to load, in ns since the epoch
-        :param time_selection: Kind of time selectoin to apply:
+        :param time_selection: Kind of time selection to apply:
         - skip: Do not select a time range, even if other arguments say so
         - touching: select things that (partially) overlap with the range
         - fully_contained: (default) select things fully contained in the range

strax/plugin.py

@@ -638,6 +638,13 @@ def do_compute(self, chunk_i=None, **kwargs):
 class LoopPlugin(Plugin):
     """Plugin that disguises multi-kind data-iteration by an event loop
     """
+    # time_selection: Kind of time selection to apply:
+    # - touching: select things that (partially) overlap with the range.
+    # The number of samples to be desired to overlapped can be set by
+    # self.touching_window. Otherwise 0 is assumed (see strax.touching_windows)
+    # - fully_contained: (default) select things fully contained in the range
+    time_selection = 'fully_contained'
+
     def compute(self, **kwargs):
         # If not otherwise specified, data kind to loop over
         # is that of the first dependency (e.g. events)
@@ -670,7 +677,17 @@ def compute(self, **kwargs):
                          for x in examples]))
 
             if k != loop_over:
-                r = strax.split_by_containment(things, base)
+                if self.time_selection == 'fully_contained':
+                    r = strax.split_by_containment(things, base)
+                elif self.time_selection == 'touching':
+                    window = 0
+                    if hasattr(self, 'touching_window'):
+                        window = self.touching_window
+                    r = strax.split_touching_windows(things,
+                                                     base,
+                                                     window=window)
+                else:
+                    raise RuntimeError('Unknown time_selection')
                 if len(r) != len(base):
                     raise RuntimeError(f"Split {k} into {len(r)}, "
                                        f"should be {len(base)}!")

strax/processing/general.py

@@ -266,6 +266,32 @@ def overlap_indices(a1, n_a, b1, n_b):
     return (a_start, a_end), (b_start, b_end)
 
 
+@export
+def split_touching_windows(things, containers, window=0):
+    """
+    Split things by their containers and return a list of length containers
+    :param things: Sorted array of interval-like data
+    :param containers: Sorted array of interval-like data
+    :param window: threshold distance for touching check
+    For example:
+       - window = 0: things must overlap one sample
+       - window = -1: things can start right after container ends
+         (i.e. container endtime equals the thing starttime, since strax
+          endtimes are exclusive)
+    :return:
+    """
+    windows = touching_windows(things, containers, window)
+    return _split_by_window(things, windows)
+
+
+@numba.njit
+def _split_by_window(r, windows):
+    result = []
+    for w in windows:
+        result.append(r[w[0]:w[1]])
+    return result
+
+
 @export
 def touching_windows(things, containers, window=0):
     """Return array of (start, exclusive end) indices into things which extend

tests/test_loop_plugins.py

@@ -36,6 +36,7 @@ def drop_random(input_array: np.ndarray) -> np.ndarray:
 def _loop_test_inner(big_data,
                      nchunks,
                      target='added_thing',
+                     time_selection='fully_contained',
                      force_value_error=False):
     """
     Test loop plugins for random data. For this test we are going to
@@ -110,6 +111,9 @@ class AddBigToSmall(strax.LoopPlugin):
         provides = 'added_thing'
         loop_over = 'big_kinda_data'  # Also just test this feature
 
+        def setup(self):
+            self.time_selection = time_selection
+
         def infer_dtype(self):
             # Get the dtype from the dependency
             return self.deps['big_thing'].dtype
@@ -221,3 +225,19 @@ def test_value_error_for_loop_plugin(big_data, nchunks):
     except ValueError:
         # Good we got the ValueError we wanted
         pass
+
+
+@given(get_some_array().filter(lambda x: len(x) >= 0),
+       strategies.integers(min_value=1, max_value=10))
+@settings(deadline=None)
+@example(
+    big_data=np.array(
+        [(0, 0, 1, 1),
+         (1, 1, 1, 1),
+         (5, 2, 2, 1),
+         (11, 4, 2, 4)],
+        dtype=full_dt_dtype),
+    nchunks=2)
+def test_loop_plugin_tw(big_data, nchunks):
+    """Test the loop plugin for random data"""
+    _loop_test_inner(big_data, nchunks, time_selection='touching')