File size: 4,560 Bytes
534cc76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Validate and clean one aligned native-rate recording for the signal encoder."""
from dataclasses import dataclass, asdict, replace
import hashlib
import json
import neurokit2 as nk
import numpy as np
import scipy
import torch
from .inputs import Channel, UnsupportedSignalError, validate_state
from .physiology import clean_channel, check_eeg_channels
from .provenance import PREPARATION_CONTRACT, NATIVE_CONTRACT
from .vendor import signal
@dataclass(frozen=True)
class PreparationConfig:
max_channels: int = signal.MAX_CHANNELS
max_samples: int = signal.MAX_SAMPLES
def __post_init__(self):
if any(isinstance(v, bool) or not isinstance(v, int) or v <= 0 for v in asdict(self).values()):
raise ValueError("Preparation budgets must be positive integers")
if self.max_channels > signal.MAX_CHANNELS:
raise ValueError(f"Encoder permits at most {signal.MAX_CHANNELS} channels")
if not signal.MIN_SAMPLES <= self.max_samples <= signal.MAX_SAMPLES:
raise ValueError(f"Encoder sample budget must be between {signal.MIN_SAMPLES} and {signal.MAX_SAMPLES}")
@dataclass
class EncoderInput:
samples: torch.Tensor # [C,L], the complete aligned recording
def to(self, device):
return EncoderInput(self.samples.to(device))
@dataclass
class PreparedState:
channels: list[Channel] # Samples retain their declared native rates
identity: str
audit: dict
class SignalPreprocessor:
"""Bypass NeuroKit for processed channels and preserve every native sampling rate."""
def __init__(self, config=None):
self.config = config or PreparationConfig()
def identity(self):
return {"contract": PREPARATION_CONTRACT, "config": asdict(self.config),
"native_processing": NATIVE_CONTRACT, "neurokit2": nk.__version__,
"numpy": np.__version__, "scipy": scipy.__version__, "signal_condition": ["raw", "processed"],
"sampling": "native-no-resampling", "format": "one-aligned-recording",
"channel_order": "as-supplied", "minimum_samples": signal.MIN_SAMPLES,
"encoder": signal.IDENTITY}
def prepare(self, state):
"""Validate all channels and run NeuroKit only on those marked raw."""
channels = validate_state(state, self.config.max_channels, self.config.max_samples)
prepared, quality, raw_channels, raw_audits = [], [], [], []
for channel in channels:
if channel.metadata["signal_condition"] == "raw":
values, observations = clean_channel(channel)
values = np.asarray(values, dtype=np.float32)
if values.shape != channel.samples.shape or not np.isfinite(values).all():
raise UnsupportedSignalError("Native cleaning changed sample count or produced nonfinite values")
current = replace(channel, samples=values)
observations["signal_condition"] = "raw"
raw_channels.append(current)
raw_audits.append(observations)
else:
current = channel
observations = {"sensor_type": channel.sensor_type, "sampling_rate_hz": channel.rate,
"signal_condition": "processed", "cleaning": "skipped"}
prepared.append(current)
quality.append(observations)
if raw_channels:
check_eeg_channels(raw_channels, [c.samples for c in raw_channels], raw_audits)
audit = {"duration_s": channels[0].metadata["duration_s"], "channels": len(prepared),
"preparation": self.identity(), "quality": quality}
digest = hashlib.sha256(json.dumps(self.identity(), sort_keys=True).encode())
for channel in prepared:
digest.update(json.dumps(channel.metadata, sort_keys=True, default=lambda x: np.asarray(x).tolist()).encode())
digest.update(channel.samples.tobytes())
return PreparedState(prepared, digest.hexdigest(), audit)
def format(self, prepared):
"""Stack the complete recording in supplied channel order without external padding."""
if prepared.audit.get("preparation") != self.identity():
raise ValueError("Prepared signal does not match current preprocessing")
samples = torch.from_numpy(np.stack([c.samples for c in prepared.channels]))
return EncoderInput(samples)
|