Motor Unit Analysis

def brace_pic_from_spike_train(
    spikes: np.ndarray,
    reference: np.ndarray,
    fsamp: float,
    *,
    kind: str = "auto",
    smoother: Optional[Callable[..., Tuple[np.ndarray, np.ndarray]]] = None,
    n_eval: Optional[int] = None,
    t_eval: Optional[np.ndarray] = None,
    ci: Union[bool, float] = False,
    ci_options: Optional[Union[CIOptions, Mapping[str, Any]]] = None,
    **brace_kwargs: Any,
) -> BracePICResult:
    """Compute brace metrics directly from a MU spike train.

    The function relies on the existing ``discharge_rate.py`` utilities in
    ``hdsemg_shared.motor_unit``; those functions are intentionally not
    duplicated here.

    Parameters
    ----------
    spikes : np.ndarray
        Binary spike train or discharge sample indices.
    reference : np.ndarray
        Reference force/torque trace sampled at ``fsamp``.
    fsamp : float
        Sampling frequency in Hz.
    kind : {"auto", "binary", "indices"}
        Interpretation of ``spikes``.
    smoother : callable, optional
        Custom smoother ``smoother(times, rate, t_eval) -> (t_eval, smooth)``.
        Defaults to ``smooth_discharge_rate_svr``.
    n_eval : int, optional
        Number of evaluation samples between recruitment and derecruitment.
        Defaults to 2048-Hz sampling over the active interval.
    t_eval : np.ndarray, optional
        Explicit evaluation times.
    ci, ci_options
        Forwarded to :func:`compute_brace_pic`.
    **brace_kwargs
        Additional arguments forwarded to :func:`compute_brace_pic`.
    """
    from .discharge_rate import (
        firing_times_from_binary,
        firing_times_from_indices,
        instantaneous_discharge_rate,
        smooth_discharge_rate_svr,
    )

    spikes = np.asarray(spikes)
    reference = np.asarray(reference, dtype=np.float64)
    if reference.ndim != 1:
        raise ValueError("reference must be a 1D array.")
    if fsamp <= 0:
        raise ValueError("fsamp must be positive.")

    if kind == "auto":
        unique = np.unique(spikes)
        is_binary = np.all(np.isin(unique, (0, 1))) and spikes.size == reference.size
        kind = "binary" if is_binary else "indices"

    if kind == "binary":
        firing_times = firing_times_from_binary(spikes, fsamp)
    elif kind == "indices":
        firing_times = firing_times_from_indices(spikes, fsamp)
    else:
        raise ValueError("kind must be one of {'auto', 'binary', 'indices'}.")

    rate_times, rate = instantaneous_discharge_rate(firing_times)

    if smoother is None:
        smoother = smooth_discharge_rate_svr

    if t_eval is None:
        if n_eval is None:
            n_eval = max(3, int(round((rate_times[-1] - rate_times[0]) * float(fsamp))) + 1)
        t_eval = np.linspace(rate_times[0], rate_times[-1], int(n_eval))
    else:
        t_eval = np.asarray(t_eval, dtype=np.float64)

    t_eval, smooth_rate = smoother(rate_times, rate, t_eval)
    ref_time = np.arange(reference.size, dtype=np.float64) / float(fsamp)
    ref_on_rate = np.interp(t_eval, ref_time, reference)

    return compute_brace_pic(
        smooth_rate,
        ref_on_rate,
        fsamp=fsamp,
        time=t_eval,
        ci=ci,
        ci_options=ci_options,
        **brace_kwargs,
    )

def compute_brace_pic(
    discharge_rate: np.ndarray,
    reference: np.ndarray,
    *,
    recruitment_idx: Optional[int] = None,
    peak_idx: Optional[int] = None,
    peak_reference_idx: Optional[int] = None,
    fsamp: Optional[float] = None,
    time: Optional[np.ndarray] = None,
    reference_unit: str = "%MVT",
    discharge_unit: str = "pps",
    distance_mode: str = "positive",
    phase_fit: str = "endpoints",
    recruitment_window: int = 1,
    peak_window: int = 1,
    brace_window: int = 1,
    peak_torque_tolerance_s: float = 0.0,
    ci: Union[bool, float] = False,
    ci_options: Optional[Union[CIOptions, Mapping[str, Any]]] = None,
    ci_method: Optional[CIMethod] = None,
    ci_interval: Optional[IntervalKind] = None,
    ci_interval_method: Optional[IntervalKind] = None,
    ci_n_draws: Optional[int] = None,
    ci_n_jobs: Optional[int] = None,
    ci_random_state: Optional[int] = None,
    random_state: Optional[int] = None,
) -> BracePICResult:
    """Compute brace-method PIC metrics for one smoothed MU discharge trace.

    Parameters
    ----------
    discharge_rate : np.ndarray
        Smoothed continuous discharge-rate trace in pps.  NaNs outside the MU's
        active period are allowed and are ignored when no explicit indices are
        supplied.
    reference : np.ndarray
        Reference force/torque trace sampled on the same time base, usually
        percent MVT/MVC.
    recruitment_idx, peak_idx, peak_reference_idx : int, optional
        Indices into the original arrays.  Defaults are: first finite active
        sample, peak discharge after recruitment, and peak reference.
    fsamp : float, optional
        Sampling frequency in Hz.  Used for time axes and tolerance conversion.
    time : np.ndarray, optional
        Explicit time axis in seconds.  If omitted and ``fsamp`` is supplied,
        ``np.arange(n) / fsamp`` is used.
    distance_mode : {"positive", "absolute"}
        ``"positive"`` selects the largest above-line deviation, which is the
        PIC-amplification interpretation.  ``"absolute"`` selects the largest
        magnitude deviation and can flag below-line curvature as a brace.
    phase_fit : {"endpoints", "ols"}
        ``"endpoints"`` uses chords recruitment->brace and brace->peak.
        ``"ols"`` fits least-squares lines to the two phases, including the
        brace sample in both phases.
    recruitment_window, peak_window, brace_window : int
        Optional local averaging windows in samples for endpoint/brace points.
        Defaults to 1 sample, matching direct geometric quantification.
    peak_torque_tolerance_s : float
        Optional tolerance for the "peak discharge after peak torque" check.
        The paper criterion is strict; default is 0.
    ci : bool or float
        ``False`` disables uncertainty estimation.  ``True`` requests the
        default 95% HDI.  A number, e.g. ``95``, requests that interval level.
        When CI is enabled, the scalar metric attributes on the returned
        result are the means of the successful uncertainty draws.  When CI is
        disabled, they are the deterministic estimates from the input trace.
    ci_options : CIOptions or dict, optional
        Detailed uncertainty options.
    ci_method, ci_interval, ci_interval_method, ci_n_draws, ci_n_jobs,
    ci_random_state, random_state : optional
        Convenience overrides for fields in ``ci_options``.

    Returns
    -------
    BracePICResult
        A structured result for one MU.  The main scalar fields are:

        ``brace_height_norm``
            The primary PIC brace metric, reported as percent right-triangle
            height (% rTri).  Larger positive values mean the discharge-rate
            trace bows farther above the recruitment-to-peak linear reference.
        ``brace_height``
            The same deviation expressed as an equivalent vertical
            discharge-rate difference in pps.
        ``acceleration_slope`` and ``attenuation_slope``
            Raw-unit slopes for the recruitment-to-brace and brace-to-peak
            phases, respectively, typically in pps/%MVT.
        ``angle``
            Reflex angle at the brace point in degrees.  A linear trace is
            180 degrees; stronger bowing increases the angle.
        ``recruitment_idx``, ``brace_idx``, ``peak_idx``
            Indices into the original input arrays identifying the analysed
            geometry points.
        ``valid`` and ``exclusion_reasons``
            Paper-style inspection flags.  Invalid results are returned rather
            than dropped, so callers can decide whether to filter them.
        ``x``, ``y``, and ``time``
            The finite recruitment-to-peak segment used for the calculation.
        ``ci``
            Optional uncertainty summary.  When CI is requested and finite
            draws are available, the scalar metric fields above are draw means;
            the deterministic geometry points and indices still describe the
            original input trace used to generate the draw set.
    """
    result = _compute_brace_pic_core(
        discharge_rate=discharge_rate,
        reference=reference,
        recruitment_idx=recruitment_idx,
        peak_idx=peak_idx,
        peak_reference_idx=peak_reference_idx,
        fsamp=fsamp,
        time=time,
        reference_unit=reference_unit,
        discharge_unit=discharge_unit,
        distance_mode=distance_mode,
        phase_fit=phase_fit,
        recruitment_window=recruitment_window,
        peak_window=peak_window,
        brace_window=brace_window,
        peak_torque_tolerance_s=peak_torque_tolerance_s,
    )

    opts = _make_ci_options(
        ci,
        ci_options=ci_options,
        method=ci_method,
        interval=ci_interval if ci_interval is not None else ci_interval_method,
        n_draws=ci_n_draws,
        n_jobs=ci_n_jobs,
        random_state=ci_random_state if ci_random_state is not None else random_state,
    )
    if opts is not None:
        result.ci = _compute_ci(
            result,
            full_reference=np.asarray(reference, dtype=np.float64),
            full_discharge=np.asarray(discharge_rate, dtype=np.float64),
            full_time=np.asarray(time, dtype=np.float64) if time is not None else None,
            fsamp=fsamp,
            opts=opts,
            core_kwargs={
                "reference_unit": reference_unit,
                "discharge_unit": discharge_unit,
                "distance_mode": distance_mode,
                "phase_fit": phase_fit,
                "recruitment_window": recruitment_window,
                "peak_window": peak_window,
                "brace_window": brace_window,
                "peak_torque_tolerance_s": peak_torque_tolerance_s,
            },
        )
        _apply_ci_point_estimates(result)

    return result

def compute_brace_pic_openhdemg_all(
    emgfile: Mapping[str, Any],
    *,
    smoothfits: Optional[Any] = None,
    ci: Union[bool, float] = False,
    ci_options: Optional[Union[CIOptions, Mapping[str, Any]]] = None,
    **brace_kwargs: Any,
) -> Tuple[Any, List[BracePICResult]]:
    """Compute brace-method PIC metrics for all MUs in an ``openhdemg`` object.

    If ``smoothfits`` is supplied, it is used directly and no SVR smoothing is
    performed.  This supports validation against externally smoothed traces such
    as manually digitized Fig. 1 discharge-rate curves.

    Parameters
    ----------
    emgfile : mapping
        ``openhdemg`` file object containing at least ``REF_SIGNAL``, ``FSAMP``,
        and either ``NUMBER_OF_MUS`` or ``MUPULSES``.
    smoothfits : pandas.DataFrame or array-like, optional
        Smoothed discharge rates with shape ``(samples, MUs)``.  NaNs outside MU
        activity are allowed.  If omitted, ``openhdemg.library.compute_svr`` is
        called.
    ci, ci_options, **brace_kwargs
        Forwarded to :func:`compute_brace_pic`.

    Returns
    -------
    summary_df : pandas.DataFrame
        One row per MU.
    results : list[BracePICResult]
        Structured per-MU results.
    """
    import pandas as pd

    fsamp = float(emgfile["FSAMP"])
    ref = _reference_array_from_openhdemg(emgfile)

    if smoothfits is None:
        import openhdemg.library as emg
        svrfits = emg.compute_svr(emgfile)
        smoothfits = pd.DataFrame(svrfits["gensvr"]).transpose()

    smooth_arr = np.asarray(smoothfits, dtype=np.float64)
    if smooth_arr.ndim == 1:
        smooth_arr = smooth_arr[:, None]
    if smooth_arr.shape[0] != ref.size and smooth_arr.shape[1] == ref.size:
        smooth_arr = smooth_arr.T
    if smooth_arr.shape[0] != ref.size:
        raise ValueError(
            f"smoothfits must have {ref.size} rows to match REF_SIGNAL; "
            f"got shape {smooth_arr.shape}."
        )

    results: List[BracePICResult] = []
    rows: List[Dict[str, Any]] = []
    for mu in range(smooth_arr.shape[1]):
        try:
            res = compute_brace_pic(
                smooth_arr[:, mu],
                ref,
                fsamp=fsamp,
                ci=ci,
                ci_options=ci_options,
                **brace_kwargs,
            )
            row = res.as_dict()
            row["mu"] = mu
        except Exception as exc:
            res = None
            row = {"mu": mu, "valid": False, "error": str(exc)}
        results.append(res)
        rows.append(row)

    return pd.DataFrame(rows), results

def plot_brace(
    result: BracePICResult,
    *,
    ax: Optional[Any] = None,
    show_ci: bool = True,
    ci_shadow: str = "sd",
    ci_metric_label: bool = True,
    show_points: bool = True,
    show_scale_bars: bool = False,
    scale_reference: float = 10.0,
    scale_discharge: float = 10.0,
    equal_scale: bool = True,
    title: Optional[str] = None,
    trace_kwargs: Optional[Mapping[str, Any]] = None,
    ci_kwargs: Optional[Mapping[str, Any]] = None,
) -> Any:
    """Plot brace geometry with optional CI mean and SD/interval shadow.

    Parameters
    ----------
    result : BracePICResult
        Output from :func:`compute_brace_pic`.
    ax : matplotlib Axes, optional
        Existing axes.  If omitted, a new figure and axes are created.
    show_ci : bool
        If true and ``result.ci`` has a trace summary, plot the CI draw mean and
        uncertainty shadow.
    ci_shadow : {"sd", "interval"}
        ``"sd"`` plots mean ± one SD.  ``"interval"`` uses the pointwise draw
        interval when draw traces are available; if unavailable, it falls back to
        SD.
    show_scale_bars : bool
        Add reference and discharge-rate scale bars.
    scale_reference, scale_discharge : float
        Scale-bar lengths used only when ``show_scale_bars=True``.
    equal_scale : bool
        Set ``ax.set_aspect(scale_reference / scale_discharge)``.  Disabled by
        default so matplotlib can fill the plot area with normal autoscaling.
    """
    import matplotlib.pyplot as plt

    if ax is None:
        _, ax = plt.subplots(figsize=(6, 5))

    trace_style = {"linewidth": 2.0, "label": "smoothed discharge"}
    if trace_kwargs:
        trace_style.update(trace_kwargs)

    x = np.asarray(result.x, dtype=float)
    y = np.asarray(result.y, dtype=float)

    if show_ci and result.ci is not None and result.ci.trace_reference is not None:
        ci_style = {"alpha": 0.18, "linewidth": 0.0, "label": "CI draw SD"}
        if ci_kwargs:
            ci_style.update(ci_kwargs)
        xr = result.ci.trace_reference
        ym = result.ci.trace_mean
        ys = result.ci.trace_sd
        if ym is not None and ys is not None:
            ax.plot(xr, ym, linewidth=1.2, alpha=0.8, label="CI draw mean")
            ax.fill_between(xr, ym - ys, ym + ys, **ci_style)

    ax.plot(x, y, **trace_style)

    # The brace-height segment is drawn to the raw perpendicular projection on
    # the recruitment-to-peak line.  This makes the plotted segment visibly
    # perpendicular when x and y are displayed with the same data scaling.
    ax.plot(
        [result.recruitment_reference, result.peak_reference],
        [result.recruitment_rate, result.peak_rate],
        linestyle="--",
        linewidth=1.0,
        label="recruitment-peak line",
    )
    ax.plot(
        [result.recruitment_reference, result.brace_reference],
        [result.recruitment_rate, result.brace_rate],
        linewidth=2.0,
        label="acceleration",
    )
    ax.plot(
        [result.brace_reference, result.peak_reference],
        [result.brace_rate, result.peak_rate],
        linewidth=2.0,
        label="attenuation",
    )
    ax.plot(
        [result.brace_reference, result.projection_reference],
        [result.brace_rate, result.projection_rate],
        linewidth=2.0,
        label="brace height",
    )

    if show_points:
        ax.plot(result.recruitment_reference, result.recruitment_rate, "o", markersize=5)
        ax.plot(result.brace_reference, result.brace_rate, "o", markersize=5)
        ax.plot(result.peak_reference, result.peak_rate, "o", markersize=5)

    if ci_metric_label and result.ci is not None:
        label = _format_ci_label(result)
        ax.text(0.02, 0.98, label, transform=ax.transAxes, va="top", ha="left",
                fontsize=9, bbox={"boxstyle": "round", "alpha": 0.12})

    ax.set_xlabel(f"Reference ({result.reference_unit})")
    ax.set_ylabel(f"Discharge rate ({result.discharge_unit})")
    if title is not None:
        ax.set_title(title)

    if equal_scale:
        ax.set_aspect(float(scale_reference) / float(scale_discharge), adjustable="datalim")

    if show_scale_bars:
        _add_scale_bars(ax, scale_reference, scale_discharge, result.reference_unit, result.discharge_unit)

    return ax

def firing_times_from_binary(spike_train: np.ndarray, fsamp: float) -> np.ndarray:
    """
    Convert a binary spike train into discharge times in seconds.

    Parameters
    ----------
    spike_train : np.ndarray
        One-dimensional binary array (non-zero where the MU discharged).
    fsamp : float
        Sampling frequency of the spike train in Hz.

    Returns
    -------
    np.ndarray
        Sorted discharge times in seconds.

    Raises
    ------
    ValueError
        If ``spike_train`` is not one-dimensional or ``fsamp`` is not positive.
    """
    spike_train = np.asarray(spike_train)
    if spike_train.ndim != 1:
        raise ValueError("spike_train must be a 1D array.")
    if fsamp <= 0:
        raise ValueError("fsamp must be a positive number.")
    indices = np.flatnonzero(spike_train)
    return indices.astype(np.float64) / float(fsamp)

def firing_times_from_indices(indices: np.ndarray, fsamp: float) -> np.ndarray:
    """
    Convert discharge sample indices into discharge times in seconds.

    Parameters
    ----------
    indices : np.ndarray
        One-dimensional array of discharge sample indices.
    fsamp : float
        Sampling frequency in Hz.

    Returns
    -------
    np.ndarray
        Sorted discharge times in seconds.

    Raises
    ------
    ValueError
        If ``indices`` is not one-dimensional or ``fsamp`` is not positive.
    """
    indices = np.asarray(indices)
    if indices.ndim != 1:
        raise ValueError("indices must be a 1D array.")
    if fsamp <= 0:
        raise ValueError("fsamp must be a positive number.")
    return np.sort(indices.astype(np.float64)) / float(fsamp)

def instantaneous_discharge_rate(
    firing_times: np.ndarray,
    min_discharges: int = MIN_DISCHARGES,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute the instantaneous discharge rate from discharge times.

    The instantaneous discharge rate is the reciprocal of the inter-spike
    interval (ISI) and is, by convention, assigned to the time of the *later*
    spike of each pair (Beauchamp et al. 2023).

    Parameters
    ----------
    firing_times : np.ndarray
        One-dimensional array of discharge times in seconds (need not be sorted).
    min_discharges : int, optional
        Minimum number of discharges required. MUs with fewer discharges are
        rejected (default: ``MIN_DISCHARGES`` = 10).

    Returns
    -------
    times : np.ndarray
        Times (s) at which each discharge-rate sample is defined (later spike of
        each ISI pair); length ``len(firing_times) - 1``.
    rate : np.ndarray
        Instantaneous discharge rate in pulses per second (pps).

    Raises
    ------
    ValueError
        If ``firing_times`` is not one-dimensional, has fewer than
        ``min_discharges`` entries, or contains non-increasing times (ISI <= 0).
    """
    firing_times = np.asarray(firing_times, dtype=np.float64)
    if firing_times.ndim != 1:
        raise ValueError("firing_times must be a 1D array.")
    if firing_times.size < min_discharges:
        raise ValueError(
            f"MU has {firing_times.size} discharges, fewer than the required "
            f"minimum of {min_discharges}; excluded from analysis."
        )
    times = np.sort(firing_times)
    isi = np.diff(times)
    if np.any(isi <= 0):
        raise ValueError("Discharge times must be strictly increasing (ISI > 0).")
    rate = 1.0 / isi
    return times[1:], rate

def smooth_discharge_rate_svr(
    times: np.ndarray,
    rate: np.ndarray,
    t_eval: Optional[np.ndarray] = None,
    *,
    C: float = 10.0,
    epsilon: float = 0.1,
    gamma="scale",
    kernel: str = "rbf",
    **svr_kwargs,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Smooth an instantaneous discharge-rate trace with Support Vector Regression.

    This reproduces the smoothing step of Beauchamp et al. (2023), who trained an
    SVR model (MATLAB ``fitrsvm``) to predict discharge rate as a function of
    time. It requires :mod:`scikit-learn`, which is *not* a hard dependency of
    ``hdsemg_shared``; it is imported lazily so that the rest of the package
    works without it.

    Parameters
    ----------
    times : np.ndarray
        Times (s) of the instantaneous discharge-rate samples.
    rate : np.ndarray
        Instantaneous discharge rate (pps) at ``times``.
    t_eval : np.ndarray, optional
        Times (s) at which to evaluate the smoothed trace. Defaults to a dense
        linear grid spanning ``[times.min(), times.max()]`` (200 points).
    C, epsilon, gamma, kernel : SVR hyperparameters
        Passed to :class:`sklearn.svm.SVR`. The defaults are reasonable starting
        points; tune them for your data as recommended by Beauchamp et al.
    **svr_kwargs
        Additional keyword arguments forwarded to :class:`sklearn.svm.SVR`.

    Returns
    -------
    t_eval : np.ndarray
        The evaluation times.
    smooth_rate : np.ndarray
        Smoothed discharge rate (pps) at ``t_eval``.

    Raises
    ------
    ImportError
        If scikit-learn is not installed.
    ValueError
        If ``times`` and ``rate`` have mismatched shapes.
    """
    try:
        from sklearn.svm import SVR
    except ImportError as exc:  # pragma: no cover - exercised only without sklearn
        raise ImportError(
            "smooth_discharge_rate_svr requires scikit-learn. "
            "Install it with `pip install scikit-learn`, or supply your own "
            "pre-smoothed discharge-rate trace to the brace-height functions."
        ) from exc

    times = np.asarray(times, dtype=np.float64)
    rate = np.asarray(rate, dtype=np.float64)
    if times.shape != rate.shape or times.ndim != 1:
        raise ValueError("times and rate must be 1D arrays of equal length.")

    if t_eval is None:
        t_eval = np.linspace(times.min(), times.max(), 200)
    t_eval = np.asarray(t_eval, dtype=np.float64)

    model = SVR(C=C, epsilon=epsilon, gamma=gamma, kernel=kernel, **svr_kwargs)
    model.fit(times.reshape(-1, 1), rate)
    smooth_rate = model.predict(t_eval.reshape(-1, 1))
    return t_eval, smooth_rate