pytesprocess 0.1.1__py3-none-any.whl

pytesprocess/process/feature_plan.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+from pytesprocess.utils import utils
+
+
+__all__ = [
+    'TraceSpec',
+    'OFModelSpec',
+    'FeatureAlgorithmPlan',
+    'FeaturePlan',
+]
+
+
+# Keep the resource declaration explicit.  Algorithms not listed here are
+# ordinary trace algorithms (including user/external extractors).
+_FFT_ALGORITHMS = {'psd_amp', 'psd_peaks', 'phase'}
+_OF_PREFIXES = ('of1x1', 'of1x2x2', 'of1x3x3', 'ofnxm', 'ofnxmx2')
+
+
+@dataclass(frozen=True, order=True)
+class TraceSpec:
+    """Raw event-window definition used by one or more feature algorithms."""
+
+    nb_samples: int
+    nb_pretrigger_samples: int
+
+    @classmethod
+    def from_config(cls, config: Mapping[str, Any]):
+        nb_samples = config.get('nb_samples')
+        nb_pre = config.get('nb_pretrigger_samples')
+        if nb_samples is None and nb_pre is None:
+            return None
+        if nb_samples is None or nb_pre is None:
+            raise ValueError(
+                'ERROR: Feature trace length and pretrigger length must both be defined.'
+            )
+        return cls(int(nb_samples), int(nb_pre))
+
+    def as_tuple(self):
+        return (self.nb_samples, self.nb_pretrigger_samples)
+
+
+@dataclass(frozen=True)
+class OFModelSpec:
+    """Identity of static QETpy OF state.
+
+    Event pretrigger is intentionally absent: it selects a raw event window but
+    does not change the template/noise optimal-filter model.
+    """
+
+    nb_samples: int
+    csd_tag: str = 'default'
+    coupling: str = 'AC'
+    ignored_frequency_peaks: tuple = ()
+    ignore_harmonics: bool = False
+
+
+@dataclass(frozen=True)
+class FeatureAlgorithmPlan:
+    channel: str
+    feature_channel: str
+    name: str
+    base_algorithm: str
+    resource_kind: str
+    trace_spec: TraceSpec | None
+    of_model_spec: OFModelSpec | None
+    params: Mapping[str, Any]
+    physical_channels: tuple[str, ...]
+    of_signal_channels: tuple[str, ...]
+
+
+class FeaturePlan:
+    """Compile normalized feature YAML into an execution/resource plan once."""
+
+    def __init__(self, processing_config, available_channels, sample_rate_hz):
+        self._processing_config = processing_config
+        self._available_channels = list(available_channels or [])
+        self._sample_rate_hz = float(sample_rate_hz)
+        self._algorithms = []
+        self._trace_channels = {}
+        self._compile()
+
+    @property
+    def algorithms(self):
+        return tuple(self._algorithms)
+
+    @property
+    def of_algorithms(self):
+        return tuple(item for item in self._algorithms
+                     if item.resource_kind == 'optimal_filter')
+
+    @property
+    def traces_config(self):
+        if not self._trace_channels:
+            return None
+        return {
+            spec.as_tuple(): list(channels)
+            for spec, channels in self._trace_channels.items()
+        }
+
+    @staticmethod
+    def _resource_kind(base_algorithm):
+        if base_algorithm in _FFT_ALGORITHMS:
+            return 'fft'
+        if base_algorithm.startswith(_OF_PREFIXES):
+            return 'optimal_filter'
+        return 'trace'
+
+    def _compile(self):
+        for channel, channel_cfg in self._processing_config.items():
+            if not isinstance(channel_cfg, dict):
+                continue
+
+            feature_channel = channel_cfg.get('feature_channel', channel)
+            physical_channels, _ = utils.split_channel_name(
+                channel, available_channels=self._available_channels
+            )
+            physical_channels = tuple(physical_channels)
+
+            for algorithm, raw_params in channel_cfg.items():
+                if not isinstance(raw_params, dict):
+                    continue
+                if not raw_params.get('run', True):
+                    continue
+
+                params = dict(raw_params)
+                base_algorithm = params.get('base_algorithm', algorithm)
+                resource_kind = self._resource_kind(base_algorithm)
+                trace_spec = TraceSpec.from_config(params)
+
+                if resource_kind in ('fft', 'optimal_filter') and trace_spec is None:
+                    raise ValueError(
+                        f'ERROR: Algorithm "{algorithm}" for channel "{channel}" '
+                        'requires an explicit trace length/pretrigger.'
+                    )
+
+                of_spec = None
+                if resource_kind == 'optimal_filter':
+                    ignored = params.get('ignored_frequency_peaks', ())
+                    if ignored is None:
+                        ignored = ()
+                    elif not isinstance(ignored, (list, tuple)):
+                        ignored = (ignored,)
+                    else:
+                        ignored = tuple(ignored)
+                    of_spec = OFModelSpec(
+                        nb_samples=trace_spec.nb_samples,
+                        csd_tag=str(params.get('csd_tag', 'default')),
+                        coupling=str(params.get('coupling', 'AC')),
+                        ignored_frequency_peaks=tuple(ignored),
+                        ignore_harmonics=bool(params.get('ignore_harmonics', False)),
+                    )
+
+                of_signal_channels = physical_channels
+                if resource_kind == 'optimal_filter':
+                    # Preserve the established OF convention: ``|`` means
+                    # multiple OF channels, while derived ``+``/``-`` names
+                    # are one composite OF signal stored under that expression.
+                    of_signal_channels, _ = utils.split_channel_name(
+                        channel, available_channels=self._available_channels,
+                        separator='|'
+                    )
+                    of_signal_channels = tuple(of_signal_channels)
+
+                plan = FeatureAlgorithmPlan(
+                    channel=channel,
+                    feature_channel=feature_channel,
+                    name=algorithm,
+                    base_algorithm=base_algorithm,
+                    resource_kind=resource_kind,
+                    trace_spec=trace_spec,
+                    of_model_spec=of_spec,
+                    params=params,
+                    physical_channels=physical_channels,
+                    of_signal_channels=of_signal_channels,
+                )
+                self._algorithms.append(plan)
+
+                if trace_spec is not None:
+                    existing = self._trace_channels.setdefault(trace_spec, [])
+                    for chan in physical_channels:
+                        if chan not in existing:
+                            existing.append(chan)

pytesprocess/process/feature_resources.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+import qetpy as qp
+
+from .feature_plan import TraceSpec, OFModelSpec, FeatureAlgorithmPlan, FeaturePlan
+
+
+__all__ = [
+    'EventTraceCache',
+    'EventFFTData',
+    'FeatureResources',
+]
+
+
+class EventFFTData:
+    """Small signal-FFT view used by non-OF frequency-domain features.
+
+    It intentionally exposes the subset of the old OFBase signal API used by
+    psd_amp/psd_peaks/phase, without constructing a QETpy OFBase.
+    """
+
+    def __init__(self, channel, trace, sample_rate_hz):
+        self._channel = channel
+        self._fs = float(sample_rate_hz)
+        self._nbins = None
+        self._freqs = None
+        self._fft = None
+        if trace is not None:
+            self._nbins = int(trace.shape[-1])
+            freqs, fft = qp.utils.fft(np.asarray(trace), self._fs, axis=-1)
+            self._freqs = np.asarray(freqs)
+            # Match QETpy OFBase.update_signal() normalization.
+            self._fft = np.asarray(fft) / self._nbins
+
+    def is_signal_stored(self, channel):
+        return channel == self._channel and self._fft is not None
+
+    def signal_fft(self, channel, squeeze_array=False):
+        if not self.is_signal_stored(channel):
+            raise ValueError(f'ERROR: FFT data are not stored for channel {channel}.')
+        out = self._fft
+        return np.squeeze(out) if squeeze_array else np.atleast_2d(out)
+
+    def fft_freqs(self):
+        return self._freqs
+
+    def nb_samples(self):
+        return self._nbins
+
+
+class EventTraceCache:
+    """Per-event lazy cache for trace windows and their FFTs."""
+
+    def __init__(self, processing_data, sample_rate_hz, weights=None):
+        self._processing_data = processing_data
+        self._sample_rate_hz = float(sample_rate_hz)
+        self._weights = weights or {}
+        self._traces = {}
+        self._ffts = {}
+
+    def _weights_for(self, channel):
+        return self._weights.get(channel)
+
+    def get_trace(self, channel, trace_spec: TraceSpec | None):
+        key = (channel, trace_spec)
+        if key not in self._traces:
+            if trace_spec is None:
+                trace = self._processing_data.get_channel_trace(
+                    channel, weights=self._weights_for(channel)
+                )
+            else:
+                trace = self._processing_data.get_channel_trace(
+                    channel,
+                    nb_samples=trace_spec.nb_samples,
+                    nb_pretrigger_samples=trace_spec.nb_pretrigger_samples,
+                    weights=self._weights_for(channel),
+                )
+            self._traces[key] = trace
+        return self._traces[key]
+
+    def get_fft_data(self, channel, trace_spec: TraceSpec):
+        key = (channel, trace_spec)
+        if key not in self._ffts:
+            trace = self.get_trace(channel, trace_spec)
+            self._ffts[key] = EventFFTData(
+                channel, trace, self._sample_rate_hz
+            )
+        return self._ffts[key]
+
+
+@dataclass
+class _OFModelState:
+    of_base: object
+    template_definitions: dict = field(default_factory=dict)
+    current_trace_spec: TraceSpec | None = None
+    loaded_channels: set = field(default_factory=set)
+
+
+class FeatureResources:
+    """Static feature resources plus lazy per-event QETpy signal loading."""
+
+    def __init__(self, feature_plan: FeaturePlan, filter_data,
+                 sample_rate_hz, available_channels=None, verbose=True):
+        self._plan = feature_plan
+        self._filter_data = filter_data
+        self._sample_rate_hz = float(sample_rate_hz)
+        self._available_channels = list(available_channels or [])
+        self._verbose = verbose
+        self._of_models = {}
+        self._prepare()
+
+    @property
+    def number_of_models(self):
+        return len(self._of_models)
+
+    def begin_event(self):
+        # Dynamic signal state is cleared once per event. Static templates,
+        # CSDs, phis, weights and p-matrices remain in the OFBase.
+        for state in self._of_models.values():
+            state.of_base.clear_signal()
+            state.current_trace_spec = None
+            state.loaded_channels.clear()
+
+    def _get_state(self, spec: OFModelSpec):
+        state = self._of_models.get(spec)
+        if state is None:
+            state = _OFModelState(
+                of_base=qp.OFBase(self._sample_rate_hz, verbose=self._verbose)
+            )
+            self._of_models[spec] = state
+        return state
+
+    def _prepare(self):
+        if not self._plan.of_algorithms:
+            return
+        if self._filter_data is None:
+            raise ValueError(
+                'ERROR: Feature optimal-filter algorithms require a filter_file.'
+            )
+
+        for plan in self._plan.of_algorithms:
+            self._prepare_algorithm(plan)
+
+    def _prepare_algorithm(self, plan: FeatureAlgorithmPlan):
+        spec = plan.of_model_spec
+        state = self._get_state(spec)
+        of = state.of_base
+        params = plan.params
+        channel = plan.channel
+
+        # Static noise model.
+        if of.csd(channel) is None:
+            csd, _freqs, metadata = self._filter_data.get_csd(
+                channel, tag=spec.csd_tag, fold=False, return_metadata=True
+            )
+            fs_filter = metadata.get('sample_rate_hz', metadata.get('sample_rate'))
+            if fs_filter is not None and float(fs_filter) != self._sample_rate_hz:
+                raise ValueError(
+                    f'ERROR: Sample rate is inconsistent between raw data '
+                    f'({self._sample_rate_hz}) and CSD ({fs_filter}) for '
+                    f'channel {channel}.'
+                )
+            if int(csd.shape[-1]) != spec.nb_samples:
+                raise ValueError(
+                    f'ERROR: Number of samples is inconsistent between raw data '
+                    f'({spec.nb_samples}) and CSD ({csd.shape[-1]}) for '
+                    f'channel {channel}, algorithm {plan.name}.'
+                )
+            ignored = list(spec.ignored_frequency_peaks) or None
+            of.set_csd(
+                channel, csd,
+                coupling=spec.coupling,
+                ignored_frequency_peaks=ignored,
+                ignore_harmonics=spec.ignore_harmonics,
+            )
+
+        template_tag = params.get('template_tag')
+        if template_tag is None:
+            raise ValueError(
+                f'ERROR: a "template_tag" in yaml file is required for '
+                f'channel {channel}, algorithm "{plan.name}"!'
+            )
+
+        template, _time, metadata = self._filter_data.get_template(
+            channel, tag=template_tag, return_metadata=True
+        )
+        if int(template.shape[-1]) != spec.nb_samples:
+            raise ValueError(
+                f'ERROR: Number of samples is inconsistent between raw data '
+                f'({spec.nb_samples}) and template ("{template_tag}") for '
+                f'channel {channel}, algorithm {plan.name}.'
+            )
+
+        template_pre = metadata.get('nb_pretrigger_samples')
+        if template_pre is None:
+            template_pre = plan.trace_spec.nb_pretrigger_samples
+        template_pre = int(template_pre)
+        integralnorm = bool(params.get('integralnorm', False))
+        template_key = (channel, str(template_tag))
+        definition = (integralnorm, template_pre)
+        old_definition = state.template_definitions.get(template_key)
+        if old_definition is not None and old_definition != definition:
+            raise ValueError(
+                f'ERROR: Conflicting definitions for template tag "{template_tag}" '
+                f'on channel {channel}: {old_definition} vs {definition}. '
+                'Use a distinct template tag for different normalization/pretrigger.'
+            )
+
+        if old_definition is None:
+            of.add_template(
+                channel, template,
+                template_tag=template_tag,
+                pretrigger_samples=template_pre,
+                integralnorm=integralnorm,
+                overwrite=False,
+            )
+            state.template_definitions[template_key] = definition
+
+        if of.phi(channel, template_tag) is None:
+            of.calc_phi(channel, template_tag)
+
+        if plan.base_algorithm == 'ofnxmx2':
+            if ('template_group_ids' not in params or 'fit_window' not in params):
+                raise ValueError(
+                    f'ERROR: "template_group_ids" and "fit_window" required '
+                    f'in yaml file for channel {channel}, algorithm {plan.name}!'
+                )
+            of.set_time_constraints(
+                channel,
+                template_group_ids=np.asarray(params['template_group_ids']),
+                fit_window=np.asarray(params['fit_window']),
+                restrict_time_flag=bool(params.get('restrict_time_flag', True)),
+                time_constraints_tag=plan.name,
+            )
+            of.calc_p_matrix(
+                channel,
+                template_tag=template_tag,
+                time_constraints_tag=plan.name,
+            )
+
+    def get_of_base(self, plan: FeatureAlgorithmPlan,
+                    event_cache: EventTraceCache):
+        if plan.resource_kind != 'optimal_filter':
+            return None
+        state = self._of_models[plan.of_model_spec]
+
+        # One static OF model can be reused for several event pretriggers.
+        # Switching TraceSpec only replaces per-event signal state.
+        if state.current_trace_spec is None:
+            state.current_trace_spec = plan.trace_spec
+        elif state.current_trace_spec != plan.trace_spec:
+            state.of_base.clear_signal()
+            state.loaded_channels.clear()
+            state.current_trace_spec = plan.trace_spec
+
+        for channel in plan.of_signal_channels:
+            if channel in state.loaded_channels:
+                continue
+            trace = event_cache.get_trace(channel, plan.trace_spec)
+            if trace is not None:
+                state.of_base.update_signal(channel, trace, calc_fft=True)
+                state.loaded_channels.add(channel)
+
+        return state.of_base