pytesprocess 0.1.1__py3-none-any.whl

pytesprocess/salting/__init__.py

@@ -0,0 +1,7 @@
+"""Salting generation and injection helpers."""
+
+from .generator import SaltGenerator
+from .injector import SaltInjector
+from .sampling import sample_pdf, sample_dm_distributions
+
+__all__ = ["SaltGenerator", "SaltInjector", "sample_pdf", "sample_dm_distributions"]

pytesprocess/salting/generator.py

@@ -0,0 +1,364 @@
+"""Salt-event metadata generation.
+
+Generation selects raw-data coordinates and records enough metadata to
+reconstruct each injected pulse later.  It intentionally does not require raw
+waveform injection and therefore remains independent of trigger/feature reads.
+"""
+
+from __future__ import annotations
+
+
+import numpy as np
+import qetpy as qp
+
+from pytesdaqx.io import AcquisitionCatalog
+from qetpy.utils import convert_channel_name_to_list, convert_channel_list_to_name
+
+from pytesprocess.utils import extract_stream_id
+from .sampling import sample_dm_distributions
+
+
+class SaltGenerator:
+    """Generate salt metadata from an acquisition and FilterData object."""
+
+    def __init__(self, filter_data, verbose=True):
+        self._filter_data = filter_data
+        self._verbose = bool(verbose)
+        self._data_path = None
+        self._streams = None
+        self._restricted = False
+        self._catalog = None
+        self._sample_rate_hz = None
+        self._positions = None
+
+    @property
+    def sample_rate_hz(self):
+        return self._sample_rate_hz
+
+    @property
+    def catalog(self):
+        return self._catalog
+
+    def set_raw_data(self, data_path, streams=None, restricted=False):
+        self._data_path = data_path
+        self._streams = streams
+        self._restricted = bool(restricted)
+
+        catalog_full = AcquisitionCatalog(data_path, verbose=self._verbose)
+        self._catalog = catalog_full.filter(
+            measurement_types="background",
+            streams=streams,
+            restricted=self._restricted,
+        )
+        selected = self._catalog.select_files_by_stream(stream_key="stream_id")
+        if not selected:
+            raise ValueError("No background streams were found for salting.")
+        self._sample_rate_hz = self._catalog.sample_rate_hz
+        return self._catalog
+
+    def get_positions(self):
+        return self._positions
+
+    def set_positions(self, dataframe):
+        self._positions = self._ensure_stream_id(dataframe)
+
+    def clear_positions(self):
+        self._positions = None
+
+    def generate_positions(self, nrandoms, *, min_separation_msec=0,
+                           edge_exclusion_msec=0, random_seed=None):
+        if self._data_path is None:
+            raise ValueError("Raw data must be set before generating salt positions.")
+        # Local import avoids a package-initialization cycle.
+        from pytesprocess.process.randoms import Randoms
+        randoms = Randoms(
+            self._data_path,
+            streams=self._streams,
+            data_type="background",
+            restricted=self._restricted,
+            verbose=False,
+        )
+        dataframe = randoms.process(
+            nrandoms=int(nrandoms),
+            min_separation_msec=float(min_separation_msec),
+            edge_exclusion_msec=float(edge_exclusion_msec),
+            partition_target_duration_s=None,
+            random_seed=random_seed,
+            lgc_save=False,
+            lgc_output=True,
+        )
+        self._positions = self._ensure_stream_id(dataframe)
+        if self._verbose:
+            print(f"INFO: {len(self._positions)} salting events randomly selected!")
+        return self._positions
+
+    def generate_metadata(
+        self,
+        channels,
+        *,
+        template_tag="default",
+        dpdi_tag=None,
+        dpdi_poles=None,
+        collection_efficiency=1.0,
+        energy_eV=None,
+        dm_pdf_file=None,
+        nsalt=100,
+        positions=None,
+        min_separation_msec=None,
+        edge_exclusion_msec=0,
+        random_seed=None,
+        salting_livetime_s=None,
+    ):
+        """Generate one salt metadata dataframe.
+
+        Exactly one of ``energy_eV`` or ``dm_pdf_file`` must be provided.
+        Fixed-energy generation is intentionally one energy at a time; an
+        outer workflow may call this repeatedly for an energy scan.
+        """
+        if (energy_eV is None) == (dm_pdf_file is None):
+            raise ValueError(
+                'Specify exactly one of "energy_eV" or "dm_pdf_file".'
+            )
+        if (dpdi_tag is None) != (dpdi_poles is None):
+            raise ValueError(
+                'Both "dpdi_tag" and "dpdi_poles" must be set or both omitted.'
+            )
+
+        channel_list = convert_channel_name_to_list(channels)
+        channel_name = convert_channel_list_to_name(channel_list)
+        efficiencies = self._normalize_efficiency(
+            collection_efficiency, len(channel_list)
+        )
+        dpdi_tags = self._normalize_per_channel(
+            dpdi_tag, len(channel_list), label="dpdi_tag"
+        )
+        dpdi_poles_list = self._normalize_per_channel(
+            dpdi_poles, len(channel_list), label="dpdi_poles"
+        )
+
+        template, time_array, template_metadata = self._filter_data.get_template(
+            channel_name, tag=template_tag, return_metadata=True
+        )
+        template_length = int(np.asarray(template).shape[-1])
+        template_pretrigger = int(
+            template_metadata.get("nb_pretrigger_samples", template_length // 2)
+        )
+
+        if energy_eV is not None:
+            if isinstance(energy_eV, (list, tuple, np.ndarray)):
+                if len(energy_eV) != 1:
+                    raise ValueError(
+                        "generate_metadata accepts one fixed energy at a time."
+                    )
+                energy_eV = energy_eV[0]
+            energies = np.full(int(nsalt), float(energy_eV), dtype=float)
+            masses = None
+            salting_type = np.asarray(
+                [f"energy_{float(energy_eV):g}_eV"] * len(energies)
+            )
+        else:
+            energies, masses = sample_dm_distributions(
+                dm_pdf_file, int(nsalt), random_seed=random_seed
+            )
+            salting_type = np.asarray(["dm_pdf"] * len(energies))
+
+        event_count = len(energies)
+        if positions is not None:
+            position_df = self._ensure_stream_id(positions)
+            self._positions = position_df
+        else:
+            if min_separation_msec is None:
+                min_separation_msec = 1000.0 * template_length / float(
+                    self._sample_rate_hz
+                )
+            position_df = self.generate_positions(
+                event_count,
+                min_separation_msec=min_separation_msec,
+                edge_exclusion_msec=edge_exclusion_msec,
+                random_seed=random_seed,
+            )
+
+        if len(position_df) != event_count:
+            raise ValueError(
+                f"Salt position count ({len(position_df)}) does not match "
+                f"generated event count ({event_count})."
+            )
+        self._positions = position_df
+
+        scales, peaks, normalizations = self._calculate_template_scales(
+            channel_list=channel_list,
+            template=template,
+            time_array=np.asarray(time_array),
+            dpdi_tags=dpdi_tags,
+            dpdi_poles_list=dpdi_poles_list,
+            energies_eV=energies,
+            efficiencies=efficiencies,
+        )
+
+        salt_dict = {
+            "salt_template_tag": np.asarray([str(template_tag)] * event_count),
+            "salt_recoil_energy_eV": energies,
+            "saltchanname": np.asarray([channel_name] * event_count),
+            "salt_channel_expression": np.asarray([channel_name] * event_count),
+            "salting_type": salting_type,
+            "salt_template_length_samples": np.full(
+                event_count, template_length, dtype=np.int64
+            ),
+            "salt_template_pretrigger_samples": np.full(
+                event_count, template_pretrigger, dtype=np.int64
+            ),
+        }
+        if masses is not None:
+            salt_dict["salt_dm_mass_MeV"] = masses
+        if salting_livetime_s is None:
+            position_columns = set(position_df.get_column_names())
+            if "randoms_livetime_s" in position_columns and len(position_df):
+                values = position_df.evaluate(
+                    "randoms_livetime_s", array_type="numpy"
+                )
+                if len(values):
+                    salting_livetime_s = float(values[0])
+        if salting_livetime_s is not None:
+            salt_dict["salting_livetime"] = np.full(
+                event_count, float(salting_livetime_s)
+            )
+        if dpdi_tag is not None:
+            if len(set(str(value) for value in dpdi_tags)) == 1 and len(set(int(value) for value in dpdi_poles_list)) == 1:
+                salt_dict["salt_dpdi_tag"] = np.asarray(
+                    [str(dpdi_tags[0])] * event_count
+                )
+                salt_dict["salt_dpdi_poles"] = np.full(
+                    event_count, int(dpdi_poles_list[0]), dtype=np.int16
+                )
+            else:
+                for index, channel in enumerate(channel_list):
+                    salt_dict[f"salt_dpdi_tag_{channel}"] = np.asarray(
+                        [str(dpdi_tags[index])] * event_count
+                    )
+                    salt_dict[f"salt_dpdi_poles_{channel}"] = np.full(
+                        event_count, int(dpdi_poles_list[index]), dtype=np.int16
+                    )
+
+        for index, channel in enumerate(channel_list):
+            scale = scales[channel]
+            # ``salt_amplitude`` is retained as a compatibility alias.  Its
+            # historical injector treated it as the multiplier applied to the
+            # stored template, which is now the explicit definition.
+            salt_dict[f"salt_template_scale_{channel}"] = scale
+            salt_dict[f"salt_amplitude_{channel}"] = scale
+            salt_dict[f"salt_peak_amplitude_amps_{channel}"] = peaks[channel]
+            # Keep historical meaning of salt_energy_eV_<chan> (recoil energy)
+            # and add the collected/channel energy explicitly.
+            salt_dict[f"salt_energy_eV_{channel}"] = energies
+            salt_dict[f"salt_collected_energy_eV_{channel}"] = (
+                energies * efficiencies[index]
+            )
+            salt_dict[f"salt_collection_efficiency_{channel}"] = np.full(
+                event_count, efficiencies[index], dtype=float
+            )
+            if normalizations[channel] is not None:
+                salt_dict[f"salt_energy_normalization_eV_per_amp_{channel}"] = (
+                    np.full(event_count, normalizations[channel], dtype=float)
+                )
+
+        output = position_df.copy()
+        for key, values in salt_dict.items():
+            output[key] = values
+        return output
+
+    def _calculate_template_scales(
+        self, *, channel_list, template, time_array, dpdi_tags,
+        dpdi_poles_list, energies_eV, efficiencies,
+    ):
+        scales = {}
+        peaks = {}
+        normalizations = {}
+        for index, channel in enumerate(channel_list):
+            temp = self._select_template_channel(template, index, len(channel_list))
+            normalization = None
+            if dpdi_tags is not None:
+                dpdi, _ = self._filter_data.get_dpdi(
+                    channel,
+                    poles=int(dpdi_poles_list[index]),
+                    tag=dpdi_tags[index],
+                )
+                dpdi = np.asarray(dpdi)
+                if dpdi.ndim > 1:
+                    dpdi = dpdi[0]
+                normalization = float(
+                    qp.get_energy_normalization(
+                        time_array, np.asarray(temp), dpdi=dpdi, lgc_ev=True
+                    )
+                )
+                if not np.isfinite(normalization) or normalization == 0:
+                    raise ValueError(
+                        f"Invalid energy normalization for salting channel {channel}."
+                    )
+                scale = energies_eV * efficiencies[index] / normalization
+            else:
+                # Without dPdI, retain the historical convention that the
+                # stored template is already normalized for multiplication by
+                # the requested energy-like amplitude.
+                scale = energies_eV * efficiencies[index]
+
+            scale = np.asarray(scale, dtype=float)
+            scales[channel] = scale
+            peaks[channel] = np.max(np.abs(np.asarray(temp))) * np.abs(scale)
+            normalizations[channel] = normalization
+        return scales, peaks, normalizations
+
+    @staticmethod
+    def _select_template_channel(template, index, nb_channels):
+        template = np.asarray(template)
+        if nb_channels == 1:
+            if template.ndim != 1:
+                return np.asarray(template).reshape(-1)
+            return template
+        # FilterData stores multi-channel templates as [nchan, ntemplate, nsamp].
+        selected = np.asarray(template[index])
+        if selected.ndim > 1:
+            selected = selected[0]
+        return selected
+
+    @staticmethod
+    def _normalize_efficiency(value, nb_channels):
+        if np.isscalar(value):
+            return np.full(nb_channels, float(value), dtype=float)
+        values = np.asarray(value, dtype=float).reshape(-1)
+        if len(values) != nb_channels:
+            raise ValueError(
+                "collection_efficiency must be scalar or have one value per channel."
+            )
+        return values
+
+    @staticmethod
+    def _normalize_per_channel(value, nb_channels, label):
+        if value is None:
+            return None
+        if isinstance(value, (list, tuple, np.ndarray)):
+            values = list(value)
+            if len(values) != nb_channels:
+                raise ValueError(
+                    f"{label} must be scalar or have one value per channel."
+                )
+            return values
+        return [value] * nb_channels
+
+    @staticmethod
+    def _ensure_stream_id(dataframe):
+        if dataframe is None:
+            raise ValueError("Salt position dataframe is None.")
+        columns = set(dataframe.get_column_names())
+        if "stream_id" in columns:
+            return dataframe
+        if "stream_number" not in columns:
+            raise ValueError(
+                "Salt position dataframe needs stream_id or stream_number."
+            )
+        source = "stream_number"
+        numbers = dataframe.evaluate(source, array_type="numpy")
+        dataframe = dataframe.copy()
+        dataframe["stream_id"] = np.asarray(
+            [extract_stream_id(int(number)) for number in numbers]
+        )
+        return dataframe

pytesprocess/salting/injector.py

@@ -0,0 +1,329 @@
+"""Raw-waveform salt injection."""
+
+from __future__ import annotations
+
+import numpy as np
+import vaex as vx
+
+from qetpy.utils import convert_channel_name_to_list
+from pytesprocess.utils import extract_stream_id
+
+
+class SaltInjector:
+    """Inject saved salt metadata into arbitrary raw-data windows."""
+
+    def __init__(self, filter_data, dataframe=None, verbose=True):
+        self._filter_data = filter_data
+        self._dataframe = None
+        self._verbose = bool(verbose)
+        self._template_cache = {}
+        if dataframe is not None:
+            self.set_dataframe(dataframe)
+
+    @property
+    def dataframe(self):
+        return self._dataframe
+
+    @property
+    def is_modern(self):
+        if self._dataframe is None:
+            return False
+        columns = set(self._dataframe.get_column_names())
+        return {"stream_id", "stream_trigger_index"}.issubset(columns)
+
+    def set_dataframe(self, dataframe):
+        if isinstance(dataframe, str):
+            dataframe = vx.open(dataframe)
+        if dataframe is None or not hasattr(dataframe, "get_column_names"):
+            raise ValueError("Unrecognized salting dataframe argument.")
+        if len(dataframe) < 1:
+            raise ValueError("No salting events found in dataframe.")
+        self._dataframe = dataframe
+
+    def inject_window(self, channels, trace, *, stream_id,
+                      trace_start_index, include_metadata=False):
+        """Inject every salt pulse overlapping one stream-global raw window."""
+        if self._dataframe is None:
+            return (trace, {}) if include_metadata else trace
+        if not self.is_modern:
+            raise ValueError(
+                "Window-based injection requires modern salting coordinates "
+                "(stream_id and stream_trigger_index)."
+            )
+
+        sid = extract_stream_id(stream_id)
+        filtered = self._dataframe[self._dataframe["stream_id"] == sid]
+        if len(filtered) == 0:
+            return (trace, {}) if include_metadata else trace
+
+        output, metadata = self._inject_rows(
+            channels, trace, filtered,
+            trigger_column="stream_trigger_index",
+            trace_start_index=int(trace_start_index),
+            stream_id=sid,
+        )
+        if include_metadata:
+            return output, metadata
+        return output
+
+    def inject_record(self, channels, trace, *, stream_id=None,
+                      stream_number=None, global_segment_number=None,
+                      include_metadata=False):
+        """Inject salts into a full HDF5/native record.
+
+        HDF5 salt dataframes use ``global_segment_number`` and
+        ``segment_trigger_index``.
+        """
+        if self._dataframe is None:
+            return (trace, {}) if include_metadata else trace
+        columns = set(self._dataframe.get_column_names())
+        filtered = self._dataframe
+
+        required = {"global_segment_number", "segment_trigger_index"}
+        if not required.issubset(columns):
+            raise ValueError(
+                "Salting dataframe needs global_segment_number and "
+                "segment_trigger_index for HDF5 record injection."
+            )
+        if global_segment_number is None:
+            raise ValueError(
+                "global_segment_number is required for HDF5 salt injection."
+            )
+
+        filtered = filtered[
+            filtered["global_segment_number"] == int(global_segment_number)
+        ]
+        if stream_id is not None and "stream_id" in columns:
+            filtered = filtered[
+                filtered["stream_id"] == extract_stream_id(stream_id)
+            ]
+        elif stream_number is not None and "stream_number" in columns:
+            filtered = filtered[
+                filtered["stream_number"] == int(stream_number)
+            ]
+        trigger_column = "segment_trigger_index"
+        sid = extract_stream_id(
+            stream_id if stream_id is not None else stream_number
+        )
+
+        if len(filtered) == 0:
+            return (trace, {}) if include_metadata else trace
+        output, metadata = self._inject_rows(
+            channels, trace, filtered,
+            trigger_column=trigger_column,
+            trace_start_index=0,
+            stream_id=sid,
+        )
+        metadata["global_segment_number"] = int(global_segment_number)
+        if include_metadata:
+            return output, metadata
+        return output
+
+    def _inject_rows(self, channels, trace, dataframe, *, trigger_column,
+                     trace_start_index, stream_id=None):
+        channel_list = self._normalize_channels(channels)
+        trace_array = np.asarray(trace).copy()
+        input_was_1d = trace_array.ndim == 1
+        if input_was_1d:
+            trace_array = trace_array.reshape(1, -1)
+        if trace_array.ndim != 2 or len(channel_list) != trace_array.shape[0]:
+            raise ValueError(
+                "Number of channels is incompatible with salting trace shape."
+            )
+
+        columns = set(dataframe.get_column_names())
+        common = {}
+        for name in (
+            trigger_column, "salt_template_tag", "salt_channel_expression",
+            "saltchanname", "salting_type", "salt_recoil_energy_eV",
+            "salt_template_pretrigger_samples", "salt_template_length_samples",
+            "dataframe_group_name", "dataframe_group_id",
+            "dataframe_group_number", "dataframe_file_index",
+        ):
+            if name in columns:
+                common[name] = self._evaluate(dataframe, name)
+
+        scale_arrays = {}
+        for channel in channel_list:
+            scale_col = f"salt_template_scale_{channel}"
+            if scale_col not in columns:
+                scale_col = f"salt_amplitude_{channel}"
+            if scale_col in columns:
+                scale_arrays[channel] = self._evaluate(dataframe, scale_col)
+
+        injected_rows = 0
+        injected_types = []
+        injected_energies = []
+        injected_group_names = []
+        injected_group_ids = []
+        injected_group_numbers = []
+        window_start = int(trace_start_index)
+        window_stop = window_start + trace_array.shape[-1]
+
+        for row_index in range(len(dataframe)):
+            expression = self._row_value(
+                common, "salt_channel_expression", row_index,
+                fallback=self._row_value(common, "saltchanname", row_index)
+            )
+            if expression is None:
+                continue
+            expression = str(expression)
+            salt_channels = convert_channel_name_to_list(expression)
+            template_tag = str(
+                self._row_value(common, "salt_template_tag", row_index,
+                                fallback="default")
+            )
+            salt_trigger = int(common[trigger_column][row_index])
+
+            # Modern salt metadata carries template support explicitly, which
+            # lets us reject non-overlapping rows without loading templates.
+            template_length = self._row_value(
+                common, "salt_template_length_samples", row_index
+            )
+            pretrigger = self._row_value(
+                common, "salt_template_pretrigger_samples", row_index
+            )
+            if (template_length is not None and not self._is_missing(template_length)
+                    and pretrigger is not None and not self._is_missing(pretrigger)):
+                template_length = int(template_length)
+                pretrigger = int(pretrigger)
+                pulse_start = salt_trigger - pretrigger
+                pulse_stop = pulse_start + template_length
+                if pulse_stop <= window_start or pulse_start >= window_stop:
+                    continue
+                template, template_meta = self._get_template(
+                    expression, template_tag
+                )
+            else:
+                template, template_meta = self._get_template(
+                    expression, template_tag
+                )
+                template_length = int(np.asarray(template).shape[-1])
+                if pretrigger is None or self._is_missing(pretrigger):
+                    pretrigger = template_meta.get(
+                        "nb_pretrigger_samples", template_length // 2
+                    )
+                pretrigger = int(pretrigger)
+                pulse_start = salt_trigger - pretrigger
+                pulse_stop = pulse_start + template_length
+
+            overlap_start = max(window_start, pulse_start)
+            overlap_stop = min(window_stop, pulse_stop)
+            if overlap_start >= overlap_stop:
+                continue
+
+            row_injected = False
+            for trace_chan_index, channel in enumerate(channel_list):
+                if channel not in salt_channels:
+                    continue
+                salt_chan_index = salt_channels.index(channel)
+                if channel not in scale_arrays:
+                    continue
+                scale = scale_arrays[channel][row_index]
+                if self._is_missing(scale):
+                    continue
+                temp = self._select_template_channel(
+                    template, salt_chan_index, len(salt_channels)
+                )
+                pulse = np.asarray(temp, dtype=float) * float(scale)
+
+                trace_lo = overlap_start - window_start
+                trace_hi = overlap_stop - window_start
+                pulse_lo = overlap_start - pulse_start
+                pulse_hi = overlap_stop - pulse_start
+                trace_array[trace_chan_index, trace_lo:trace_hi] += pulse[pulse_lo:pulse_hi]
+                row_injected = True
+
+            if row_injected:
+                injected_rows += 1
+                value = self._row_value(common, "salting_type", row_index)
+                if value is not None and not self._is_missing(value):
+                    injected_types.append(str(value))
+                energy = self._row_value(common, "salt_recoil_energy_eV", row_index)
+                if energy is not None and not self._is_missing(energy):
+                    injected_energies.append(float(energy))
+                group_name = self._row_value(common, "dataframe_group_name", row_index)
+                group_id = self._row_value(common, "dataframe_group_id", row_index)
+                group_number = self._row_value(common, "dataframe_group_number", row_index)
+                if group_name is not None and not self._is_missing(group_name):
+                    injected_group_names.append(str(group_name))
+                if group_id is not None and not self._is_missing(group_id):
+                    injected_group_ids.append(str(group_id))
+                if group_number is not None and not self._is_missing(group_number):
+                    injected_group_numbers.append(int(group_number))
+
+        unique_types = list(dict.fromkeys(injected_types))
+        metadata = {}
+        if injected_rows:
+            metadata = {
+                "salting_type": (unique_types[0] if len(unique_types) == 1 else "multiple"),
+                "salting_types": unique_types,
+                "nb_salts": injected_rows,
+            }
+            if stream_id is not None:
+                metadata["stream_id"] = stream_id
+            names = list(dict.fromkeys(injected_group_names))
+            ids = list(dict.fromkeys(injected_group_ids))
+            numbers = list(dict.fromkeys(injected_group_numbers))
+            if len(names) == 1:
+                metadata["salting_dataframe_group_name"] = names[0]
+            if len(ids) == 1:
+                metadata["salting_dataframe_group_id"] = ids[0]
+            if len(numbers) == 1:
+                metadata["salting_dataframe_group_number"] = numbers[0]
+            if injected_energies:
+                metadata["salt_recoil_energy_eV"] = injected_energies
+
+        output = trace_array[0] if input_was_1d else trace_array
+        return output, metadata
+
+    def _get_template(self, expression, tag):
+        key = (str(expression), str(tag))
+        cached = self._template_cache.get(key)
+        if cached is not None:
+            return cached
+        template, _, metadata = self._filter_data.get_template(
+            expression, tag=tag, return_metadata=True
+        )
+        value = (template, metadata)
+        self._template_cache[key] = value
+        return value
+
+    @staticmethod
+    def _normalize_channels(channels):
+        if isinstance(channels, str):
+            return convert_channel_name_to_list(channels)
+        return list(channels)
+
+    @staticmethod
+    def _evaluate(dataframe, column):
+        values = dataframe.evaluate(column, array_type="numpy")
+        if np.ma.isMaskedArray(values):
+            values = values.filled(np.nan)
+        return np.asarray(values)
+
+    @staticmethod
+    def _row_value(common, name, index, fallback=None):
+        values = common.get(name)
+        if values is None:
+            return fallback
+        return values[index]
+
+    @staticmethod
+    def _is_missing(value):
+        if value is None:
+            return True
+        try:
+            return bool(np.isnan(value))
+        except (TypeError, ValueError):
+            return False
+
+    @staticmethod
+    def _select_template_channel(template, index, nb_channels):
+        template = np.asarray(template)
+        if nb_channels == 1:
+            return template.reshape(-1) if template.ndim != 1 else template
+        selected = np.asarray(template[index])
+        if selected.ndim > 1:
+            selected = selected[0]
+        return selected