Bug fix in pedestal estimate.

pes_to_spec/model.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Union, Tuple, Literal
 
 import sys
 import joblib
+import dask.array as da
 
 import numpy as np
 import scipy
@@ -315,6 +316,9 @@ class SelectRelevantLowResolution(TransformerMixin, BaseEstimator):
                            for delta in pulse_spacing[channel]]
                  for channel, item in X.items()
                    if channel in self.channels}
+            # convert to Numpy if it is a Dask array
+            y = {ch: [item.compute() if isinstance(item, da.Array) else item for item in v]
+                 for ch, v in y.items()}
             # pad it with zeros, if we reach the edge of the array
             for channel in y.keys():
                 y[channel] = [np.pad(y[channel][j], ((0, 0), (0, 2*self.delta_tof - y[channel][j].shape[1])))
@@ -340,6 +344,9 @@ class SelectRelevantLowResolution(TransformerMixin, BaseEstimator):
         """
         # reduce on channel and on train ID
         sum_low_res = np.mean(sum([-(v - self.pedestal[ch]) for ch, v in X.items()]), axis=0)
+        # convert to Numpy if it is a Dask array
+        if isinstance(sum_low_res, da.Array):
+            sum_low_res = sum_low_res.compute()
         axis = np.arange(0.0, sum_low_res.shape[0], 1.0)
         #widths = np.arange(10, 50, step=5)
         #peak_idx = find_peaks_cwt(sum_low_res, widths)
@@ -371,8 +378,11 @@ class SelectRelevantLowResolution(TransformerMixin, BaseEstimator):
         Returns: The object itself.
         """
         # estimate pedestal
-        X_pedestal = {ch: np.mean(v[:10])
+        self.pedestal = {ch: np.mean(v[:10])
                       for ch, v in X.items()}
+        # convert to Numpy if it is a Dask array
+        self.pedestal = {ch: v.compute() if isinstance(v, da.Array) else v
+                      for ch, v in self.pedestal.items()}
         self.tof_start = self.estimate_prompt_peak(X)
         X_tr = self.transform(X, keep_dictionary_structure=True)
         self.mean = {ch: np.mean(X_tr[ch], axis=0, keepdims=True)