From 81cf00b791b7703b7979f28d69bd606946c62fb1 Mon Sep 17 00:00:00 2001
From: Andrew Laing <a.b.r.laing@gmail.com>
Date: Thu, 26 Nov 2020 16:43:45 +0100
Subject: [PATCH] Allow pd.Series and np.ndarray as arguments in
 bufffer_function

---
 invisible_cities/detsim/buffer_functions.py | 38 +++++++++++++--------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/invisible_cities/detsim/buffer_functions.py b/invisible_cities/detsim/buffer_functions.py
index 18cab781aa..0653856a42 100644
--- a/invisible_cities/detsim/buffer_functions.py
+++ b/invisible_cities/detsim/buffer_functions.py
@@ -1,9 +1,10 @@
 import numpy  as np
 import pandas as pd
 
-from typing import  Callable
-from typing import      List
-from typing import     Tuple
+from typing import Callable
+from typing import     List
+from typing import    Tuple
+from typing import    Union
 
 from .. reco.peak_functions import indices_and_wf_above_threshold
 from .. reco.peak_functions import                 split_in_peaks
@@ -41,14 +42,17 @@ def bin_sensors(sensors   : pd.DataFrame,
 
 ## !! to-do: clarify for non-pmt versions of next
 ## !! to-do: Check on integral instead of only threshold?
-def find_signal_start(wfs          : pd.Series,
-                      bin_threshold: float    ,
-                      stand_off    : int      ) -> List[int]:
+def find_signal_start(wfs          : Union[pd.Series, np.ndarray],
+                      bin_threshold: float                       ,
+                      stand_off    : int                         ) -> List[int]:
     """
     Finds signal in the binned waveforms and
     identifies candidate triggers.
     """
-    eng_sum = wfs.sum()
+    if isinstance(wfs, np.ndarray):
+        eng_sum = wfs.sum(axis=0)
+    else:
+        eng_sum = wfs.sum()
     indices = indices_and_wf_above_threshold(eng_sum,
                                              bin_threshold).indices
     ## Just using this and the stand_off for now
@@ -122,20 +126,24 @@ def generate_slice(trigger    : int       ,
                 pad_safe(sipm_charge[:, sipm_slice], sipm_pad))
 
 
-    def position_signal(triggers   :       List,
-                        pmt_bins   : np.ndarray,
-                        pmt_charge :  pd.Series,
-                        sipm_bins  : np.ndarray,
-                        sipm_charge:  pd.Series
+    def position_signal(triggers   :       List                  ,
+                        pmt_bins   : np.ndarray                  ,
+                        pmt_charge : Union[pd.Series, np.ndarray],
+                        sipm_bins  : np.ndarray                  ,
+                        sipm_charge: Union[pd.Series, np.ndarray]
                         ) -> List[Tuple[np.ndarray, np.ndarray]]:
         """
         Synchronises the SiPMs and PMTs for each identified
         trigger and calls the padding function to fill with
         zeros where necessary.
         """
-        pmt_q  = np.asarray(pmt_charge.tolist())
-        sipm_q = np.empty((0,0))\
-          if sipm_charge.empty else np.asarray(sipm_charge.tolist())
+        if isinstance(pmt_charge, pd.Series):
+            pmt_q  = np.asarray(pmt_charge.tolist())
+            sipm_q = np.empty((0,0))\
+            if sipm_charge.empty else np.asarray(sipm_charge.tolist())
+        else:
+            pmt_q  =  pmt_charge
+            sipm_q = sipm_charge
         return [generate_slice(trigger, pmt_bins, pmt_q, sipm_bins, sipm_q)
                 for trigger in triggers                                    ]
     return position_signal