orisys-sdk/src/orisys/signal_processing.py

"""Signal helpers for force filtering and force calibration.

Batch ``lowpass_filter`` is the zero-phase Butterworth used by example charts.
``LowpassFilter`` is a causal, stateful filter for per-frame ``Sensor`` outputs.

Force calibration resolution order (when register support arrives):
  1. Sensor register / EEPROM calibration blob
  2. ``sensors/<name>/force_calibration.json``
  3. Hardcoded linear defaults below
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.signal import butter, filtfilt, sosfilt, sosfilt_zi

# Linear calibration from raw normal-force units to Newtons (ORY-BOX / ORY-MINI).
# Prefer sensors/<name>/force_calibration.json when available.
NORMAL_FORCE_CAL_SLOPE = 3.92685554668795e-05
NORMAL_FORCE_CAL_OFFSET = 0.0217596967278535

DEFAULT_SAMPLE_RATE_HZ = 30.0
DEFAULT_LOWPASS_ORDER = 2


def _parse_calibration_point(point) -> tuple[float, float]:
    if isinstance(point, dict):
        raw = point.get("fnormal", point.get("raw", point.get("x")))
        force_n = point.get("force_n", point.get("n", point.get("y")))
        if raw is None or force_n is None:
            raise ValueError(f"invalid calibration point: {point!r}")
        return float(raw), float(force_n)
    if isinstance(point, (list, tuple)) and len(point) >= 2:
        return float(point[0]), float(point[1])
    raise ValueError(f"invalid calibration point: {point!r}")


@dataclass(frozen=True)
class ForceChannelCalibration:
    """One force channel: linear coeffs and/or piecewise-linear breakpoints."""

    method: str = "linear"
    slope: float | None = None
    offset: float | None = None
    points: tuple[tuple[float, float], ...] = ()
    extrapolate: bool = True
    clamp_output_min: float | None = 0.0
    input_mode: str = "magnitude"
    signed: bool = False
    unit: str = "N"

    @classmethod
    def linear(cls, slope: float, offset: float, *, unit: str = "N") -> ForceChannelCalibration:
        return cls(method="linear", slope=float(slope), offset=float(offset), unit=unit)

    @classmethod
    def from_channel_dict(
        cls,
        block: dict,
        *,
        unit: str = "N",
        default_method: str = "piecewise_linear",
    ) -> ForceChannelCalibration | None:
        if not isinstance(block, dict):
            return None

        method = str(block.get("method", block.get("model", default_method))).lower()
        if method in ("linear", "affine"):
            slope = block.get("slope")
            offset = block.get("offset")
            if slope is None or offset is None:
                return None
            return cls.linear(float(slope), float(offset), unit=str(block.get("unit", unit)))

        if method != "piecewise_linear":
            return None

        raw_points = block.get("points", [])
        if not isinstance(raw_points, (list, tuple)) or len(raw_points) < 2:
            return None
        try:
            points = tuple(sorted(_parse_calibration_point(p) for p in raw_points))
        except (TypeError, ValueError):
            return None
        if len(points) < 2:
            return None

        return cls(
            method="piecewise_linear",
            points=points,
            extrapolate=bool(block.get("extrapolate", True)),
            clamp_output_min=block.get("clamp_output_min", 0.0),
            input_mode=str(block.get("input", "magnitude")).lower(),
            signed=bool(block.get("signed", False)),
            unit=str(block.get("unit", unit)),
        )

    def convert(self, raw: float) -> float:
        raw = float(raw)
        if self.method == "linear":
            slope = NORMAL_FORCE_CAL_SLOPE if self.slope is None else float(self.slope)
            offset = NORMAL_FORCE_CAL_OFFSET if self.offset is None else float(self.offset)
            return slope * raw + offset

        points = self.points
        if len(points) < 2:
            return raw

        if raw <= points[0][0]:
            segment = (points[0], points[1])
            if not self.extrapolate:
                return self._clamp(points[0][1])
        elif raw >= points[-1][0]:
            segment = (points[-2], points[-1])
            if not self.extrapolate:
                return self._clamp(points[-1][1])
        else:
            segment = None
            for left, right in zip(points, points[1:]):
                if left[0] <= raw <= right[0]:
                    segment = (left, right)
                    break
            if segment is None:
                return self._clamp(points[-1][1])

        (raw0, n0), (raw1, n1) = segment
        if raw1 == raw0:
            value = n0
        else:
            ratio = (raw - raw0) / (raw1 - raw0)
            value = n0 + ratio * (n1 - n0)
        return self._clamp(value)

    def convert_component(self, raw: float) -> float:
        if not self.signed:
            return self.convert(raw)
        raw = float(raw)
        sign = -1.0 if raw < 0.0 else 1.0
        return sign * self.convert(abs(raw))

    def _clamp(self, value: float) -> float:
        if self.clamp_output_min is not None:
            return max(float(self.clamp_output_min), float(value))
        return float(value)


_DEFAULT_NORMAL = ForceChannelCalibration.linear(
    NORMAL_FORCE_CAL_SLOPE,
    NORMAL_FORCE_CAL_OFFSET,
)


def _load_force_calibration_data(sensor_name: str | None) -> dict | None:
    if not sensor_name:
        return None
    try:
        from .configs import load_force_calibration

        data = load_force_calibration(sensor_name)
    except Exception:
        return None
    return data if isinstance(data, dict) else None


def _channel_from_data(data: dict | None, channel: str) -> ForceChannelCalibration | None:
    if not isinstance(data, dict):
        return None

    unit = str(data.get("unit", "N"))
    channels = data.get("channels")
    if isinstance(channels, dict):
        block = channels.get(channel)
        cal = ForceChannelCalibration.from_channel_dict(block, unit=unit)
        if cal is not None:
            return cal

    # Legacy single-channel file: {"normal_force": {"model": "linear", ...}}
    if channel == "normal":
        block = data.get("normal_force", data)
        cal = ForceChannelCalibration.from_channel_dict(
            block, unit=unit, default_method="linear"
        )
        if cal is not None:
            return cal
        if isinstance(block, dict):
            slope = block.get("slope")
            offset = block.get("offset")
            if slope is not None and offset is not None:
                return ForceChannelCalibration.linear(
                    float(slope),
                    float(offset),
                    unit=str(block.get("unit", unit)),
                )
    return None


def _normal_calibration(sensor_name: str | None = None) -> ForceChannelCalibration:
    cal = _channel_from_data(_load_force_calibration_data(sensor_name), "normal")
    return cal if cal is not None else _DEFAULT_NORMAL


def _shear_calibration(sensor_name: str | None = None) -> ForceChannelCalibration | None:
    return _channel_from_data(_load_force_calibration_data(sensor_name), "shear")


def _load_normal_force_coeffs(sensor_name: str | None = None) -> tuple[float, float]:
    """Legacy helper: linear slope/offset only (piecewise returns hardcoded defaults)."""
    cal = _normal_calibration(sensor_name)
    if cal.method == "linear" and cal.slope is not None and cal.offset is not None:
        return float(cal.slope), float(cal.offset)
    return NORMAL_FORCE_CAL_SLOPE, NORMAL_FORCE_CAL_OFFSET


def calibrate_normal_force(raw_value: float, *, sensor_name: str | None = None) -> float:
    """Convert a raw normal-force reading to Newtons.

    Uses ``sensors/<name>/force_calibration.json`` when present (piecewise or linear).
    Falls back to hardcoded linear defaults when missing.
    """
    return _normal_calibration(sensor_name).convert(raw_value)


def calibrate_shear_force(raw_value: float, *, sensor_name: str | None = None) -> float:
    """Convert a raw shear magnitude to Newtons using the shear channel curve.

    Returns ``float(raw_value)`` unchanged when no shear calibration is available.
    """
    cal = _shear_calibration(sensor_name)
    if cal is None:
        return float(raw_value)
    return cal.convert(raw_value)


def calibrate_shear_components(
    shear_x: float,
    shear_y: float,
    *,
    sensor_name: str | None = None,
) -> tuple[float, float]:
    """Convert raw shear components to Newtons.

    When the shear channel uses ``input: components`` and ``signed: true``, each
    axis is calibrated on absolute magnitude then re-signed. Otherwise the
    magnitude is calibrated and components are scaled proportionally.
    """
    sx = float(shear_x)
    sy = float(shear_y)
    cal = _shear_calibration(sensor_name)
    if cal is None:
        return sx, sy

    if cal.input_mode == "components":
        return cal.convert_component(sx), cal.convert_component(sy)

    magnitude = (sx * sx + sy * sy) ** 0.5
    if magnitude <= 0.0:
        return 0.0, 0.0
    calibrated = cal.convert(magnitude)
    scale = calibrated / magnitude
    return sx * scale, sy * scale


def _nyquist_hz(sample_rate_hz: float) -> float:
    return 0.5 * float(sample_rate_hz)


def _validate_cutoff(cutoff_hz: float, sample_rate_hz: float) -> None:
    nyquist = _nyquist_hz(sample_rate_hz)
    if cutoff_hz < 0 or cutoff_hz >= nyquist:
        raise ValueError(
            f"cutoff_hz must be in (0, {nyquist}) for sample_rate_hz={sample_rate_hz}"
        )


def lowpass_filter(
    data,
    cutoff_hz: float = 5.0,
    sample_rate_hz: float = DEFAULT_SAMPLE_RATE_HZ,
    order: int = DEFAULT_LOWPASS_ORDER,
) -> np.ndarray:
    """Zero-phase Butterworth low-pass for uniformly sampled 1-D data.

    Matches the example/viewer chart filter: ``cutoff_hz == 0`` is a no-op.
    """
    y = np.asarray(data, dtype=np.float64)
    if cutoff_hz == 0:
        return y.copy()
    if y.size == 0:
        return y

    _validate_cutoff(cutoff_hz, sample_rate_hz)
    b, a = butter(int(order), cutoff_hz / _nyquist_hz(sample_rate_hz), btype="low")
    padlen = 3 * (max(len(a), len(b)) - 1)
    if y.size <= padlen:
        return y.copy()
    try:
        return filtfilt(b, a, y)
    except ValueError:
        return y.copy()


class LowpassFilter:
    """Causal Butterworth low-pass with persistent SOS state (one sample at a time)."""

    def __init__(
        self,
        cutoff_hz: float = 0.0,
        sample_rate_hz: float = DEFAULT_SAMPLE_RATE_HZ,
        order: int = DEFAULT_LOWPASS_ORDER,
        n_channels: int = 3,
    ):
        self.n_channels = int(n_channels)
        self._sos = None
        self._zi = None
        self.configure(
            cutoff_hz=cutoff_hz,
            sample_rate_hz=sample_rate_hz,
            order=order,
        )

    @property
    def enabled(self) -> bool:
        return self._sos is not None

    def configure(
        self,
        *,
        cutoff_hz: float | None = None,
        sample_rate_hz: float | None = None,
        order: int | None = None,
    ) -> None:
        if cutoff_hz is not None:
            self.cutoff_hz = float(cutoff_hz)
        if sample_rate_hz is not None:
            self.sample_rate_hz = float(sample_rate_hz)
        if order is not None:
            self.order = int(order)

        if self.cutoff_hz <= 0:
            self._sos = None
            self._zi = None
            return

        _validate_cutoff(self.cutoff_hz, self.sample_rate_hz)
        self._sos = butter(
            self.order,
            self.cutoff_hz / _nyquist_hz(self.sample_rate_hz),
            btype="low",
            output="sos",
        )
        self._zi = None

    def reset(self, x0=None) -> None:
        """Clear filter memory. Optional ``x0`` initializes to a steady-state DC."""
        if self._sos is None:
            self._zi = None
            return
        if x0 is None:
            self._zi = None
            return
        x = np.asarray(x0, dtype=np.float64).reshape(-1)
        if x.size != self.n_channels:
            raise ValueError(f"x0 must have {self.n_channels} channels, got {x.size}")
        zi0 = sosfilt_zi(self._sos)
        self._zi = zi0[:, :, None] * x[None, None, :]

    def process(self, values) -> np.ndarray:
        """Filter one multi-channel sample. Returns float64 array of shape (n_channels,)."""
        x = np.asarray(values, dtype=np.float64).reshape(-1)
        if x.size != self.n_channels:
            raise ValueError(f"expected {self.n_channels} channels, got {x.size}")
        if self._sos is None:
            return x.copy()
        if self._zi is None:
            self.reset(x)
        y, self._zi = sosfilt(self._sos, x[None, :], zi=self._zi, axis=0)
        return np.asarray(y[0], dtype=np.float64)