src/fpgm/motion/profile.py

"""Trapezoidal (bang-coast-bang) velocity profile on a normalized path parameter.

Pure scalar math -- no robot, no URDF, no ``ArmKinematics``. Given a segment's
Cartesian translation length ``L`` and rotation angle ``Theta`` (both computed
by the caller -- see :mod:`fpgm.motion.trajectory`), :func:`trapezoidal_profile`
produces a schedule for a single shared parameter ``s in [0, 1]`` that both the
translation and rotation are driven by (``p(s) = p0 + s*(p1-p0)``,
``R(s) = R0 @ rotvec_to_matrix(s * log(R0^T R1))``). Driving both from one
scalar ``s`` is what makes translation and rotation start and stop together --
two independently-profiled axes can each individually obey their own
velocity/acceleration limits while still drifting out of sync with each other
mid-segment (one finishes early, coasts, waits for the other), which looks
wrong and complicates the velocity-clamp retiming in
:mod:`fpgm.motion.trajectory`. A single ``s`` makes that structurally
impossible.

The per-axis limits are combined into an effective limit on ``s`` itself::

    s_dot_max  = min(v_max / L,  omega_max / Theta)
    s_ddot_max = min(a_max / L,  alpha_max / Theta)

(only the terms whose denominator is nonzero participate -- see the
pure-rotation/pure-translation handling below), then the classic two-case
closed-form point-to-point (start and end at rest) trapezoidal profile is
built on ``s``:

* **Triangular** (``s_dot_max**2 / s_ddot_max >= 1``): cruise speed is never
  reached -- the segment is short enough that pure acceleration-then-
  deceleration alone covers the unit distance before ``s_dot_max`` is hit.
  ``t_acc = sqrt(1 / s_ddot_max)``, ``T = 2 * t_acc``.
* **Trapezoidal**: ``t_acc = s_dot_max / s_ddot_max``, cruise at ``s_dot_max``
  in between, ``T = 1/s_dot_max + s_dot_max/s_ddot_max``.

**Grid snapping.** DROID's trajectories -- and everything downstream that
consumes this module's output (:class:`fpgm.objects.interaction.InteractionModel`,
every renderer) -- run at a fixed 15 Hz tick. If each segment's duration were
left at its raw closed-form ``T``, consecutive segments would land on
unrelated, non-aligned time grids, and concatenating them would require either
resampling (introducing interpolation error into a path that was already
exact) or accepting timestamp drift at every segment boundary. Instead, every
segment is snapped to the nearest 15 Hz tick *at or after* its raw duration --
``n = ceil(T * 15)`` -- and then **stretched** to exactly ``T' = n / 15`` by
scaling ``s_dot_max`` and ``s_ddot_max`` down by ``k = T/T'`` and ``k**2``
respectively (``k <= 1`` always, since ``ceil`` only ever rounds up). Scaling
a velocity-and-acceleration-bounded profile uniformly down like this can only
ever make it *slower* everywhere -- it can never push ``s_dot`` or ``s_ddot``
back over the original limit -- so stretching is always safe. This is also
exactly the mechanism :class:`fpgm.motion.trajectory.TrajectoryPlanner` reuses
to re-time a segment whose realised joint velocity turned out to exceed a
joint's limit: force a larger ``min_ticks`` and the same stretch argument
applies.

Degenerate cases (handled explicitly, not by accident of the algebra):

* **Zero-length** (``L`` and ``Theta`` both ~0): nothing moves, so there is no
  meaningful ``s_dot_max``/``s_ddot_max`` to compute (both ratios above would
  divide by zero) -- a single grid tick is returned with ``s`` jumping
  straight to 1.0 and zero velocity/acceleration throughout.
* **Pure rotation** (``L ~ 0``, ``Theta > 0``): the ``v_max/L`` and
  ``a_max/L`` terms are dropped from the ``min(...)`` rather than evaluated
  (which would be a division by zero), leaving ``s_dot_max = omega_max/Theta``.
* **Pure translation** (``Theta ~ 0``, ``L > 0``): symmetric, drops the
  rotation terms.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Below this, a length (metres) or angle (radians) is treated as exactly
#: zero -- both are physical quantities from real FK/target poses, so this is
#: comfortably above float64 noise but far below any deliberate motion.
_EPS = 1e-9

#: Slack applied to the boundary-value overrides (see the end of
#: :func:`trapezoidal_profile`) -- not a tunable, just documents that s(0)/s(T)
#: are forced exact rather than left to float roundoff from the phase-boundary
#: arithmetic.
_BOUNDARY_EXACT = True


@dataclass
class ProfileConfig:
    """Named limits for :func:`trapezoidal_profile` -- never literals in the algorithm.

    ``v_max`` (the per-segment desired Cartesian speed) is deliberately *not*
    here: it varies per :class:`~fpgm.motion.types.CartesianWaypoint` (its
    ``speed_mps``), whereas these are physical limits of the arm/motion
    policy that stay fixed across a whole plan.

    Attributes:
        a_max: Max Cartesian linear acceleration, m/s^2.
        omega_max: Max Cartesian angular speed, rad/s.
        alpha_max: Max Cartesian angular acceleration, rad/s^2.
        rate_hz: The grid every segment is snapped to (see the module
            docstring) -- DROID's trajectory rate.
    """

    a_max: float = 0.5
    omega_max: float = 0.8
    alpha_max: float = 2.0
    rate_hz: float = 15.0


@dataclass
class ScalarProfile:
    """A grid-snapped trapezoidal schedule for one shared path parameter ``s``.

    Attributes:
        timestamps: ``(n_ticks + 1,)`` seconds, ``0, dt, 2*dt, ..., duration``
            -- includes the ``t=0`` boundary sample (callers stitching
            multiple segments together drop it for every segment but the very
            first, since it duplicates the previous segment's last sample).
        s, s_dot, s_ddot: ``(n_ticks + 1,)`` each, sampled at ``timestamps``.
        duration: ``timestamps[-1]``, the grid-snapped (and possibly
            retiming-stretched) total duration.
        is_triangular: Which closed-form branch this profile took.
        s_dot_max, s_ddot_max: The *stretched* (grid-snap-adjusted) limits
            actually realised -- always ``<=`` the raw, unsnapped values
            passed in, per the module docstring's stretch argument.
    """

    timestamps: np.ndarray
    s: np.ndarray
    s_dot: np.ndarray
    s_ddot: np.ndarray
    duration: float
    is_triangular: bool
    s_dot_max: float
    s_ddot_max: float

    @property
    def n_ticks(self) -> int:
        """Number of *new* grid samples (excludes the shared ``t=0`` boundary)."""
        return int(self.timestamps.shape[0]) - 1


def _sample_scurve(
    t: np.ndarray, duration: float, t_acc: float, s_dot_max: float, s_ddot_max: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Evaluate the accel/cruise/decel piecewise closed form at times ``t``.

    Works for both branches: a triangular profile is just a trapezoidal one
    whose cruise phase (``t_acc <= t <= duration - t_acc``) has zero width,
    which happens automatically when ``t_acc == duration / 2``.
    """
    s = np.empty_like(t)
    s_dot = np.empty_like(t)
    s_ddot = np.empty_like(t)

    accel = t <= t_acc
    decel = t > (duration - t_acc)
    cruise = ~accel & ~decel

    s[accel] = 0.5 * s_ddot_max * t[accel] ** 2
    s_dot[accel] = s_ddot_max * t[accel]
    s_ddot[accel] = s_ddot_max

    s_at_acc_end = 0.5 * s_ddot_max * t_acc**2
    t_cruise = t[cruise]
    s[cruise] = s_at_acc_end + s_dot_max * (t_cruise - t_acc)
    s_dot[cruise] = s_dot_max
    s_ddot[cruise] = 0.0

    t_from_end = duration - t[decel]
    s[decel] = 1.0 - 0.5 * s_ddot_max * t_from_end**2
    s_dot[decel] = s_ddot_max * t_from_end
    s_ddot[decel] = -s_ddot_max

    return s, s_dot, s_ddot


def trapezoidal_profile(
    length_m: float,
    angle_rad: float,
    v_max: float,
    cfg: ProfileConfig | None = None,
    min_ticks: int = 1,
) -> ScalarProfile:
    """Build a grid-snapped trapezoidal profile for a segment of given size.

    Args:
        length_m: Cartesian translation distance for this segment, ``||p1-p0||``.
        angle_rad: Cartesian rotation angle, ``||log(R0^T R1)||``.
        v_max: Desired peak Cartesian linear speed for this segment (e.g. a
            :class:`~fpgm.motion.types.CartesianWaypoint`'s ``speed_mps``).
        cfg: Acceleration/angular-rate limits and the grid rate; see
            :class:`ProfileConfig`.
        min_ticks: Force at least this many grid ticks even if the raw
            duration would snap to fewer -- used by
            :class:`fpgm.motion.trajectory.TrajectoryPlanner` to re-time a
            segment whose realised joint velocity exceeded a limit (see the
            module docstring's stretch argument for why forcing more ticks
            is always a safe, limit-respecting way to slow a segment down).

    Returns:
        A :class:`ScalarProfile` sampled on the grid ``0, 1/rate_hz, ...,
        duration``.
    """
    cfg = cfg if cfg is not None else ProfileConfig()
    length_m = float(length_m)
    angle_rad = float(angle_rad)
    rate_hz = float(cfg.rate_hz)
    dt = 1.0 / rate_hz
    min_ticks = max(1, int(min_ticks))

    zero_length = length_m < _EPS
    zero_angle = angle_rad < _EPS

    if zero_length and zero_angle:
        # Nothing moves: L and Theta are both ~0, so v_max/L and omega_max/Theta
        # are both divisions by zero -- there is no rate to compute at all, so
        # this is handled before any of that arithmetic runs. One grid tick
        # (the minimum a segment can occupy) is emitted with s jumping
        # straight to "arrived"; since p(s) and R(s) do not depend on s at all
        # when p0==p1 and R0==R1, *which* s value is reported never matters
        # physically -- only that it is finite and the array shapes are the
        # usual ones so callers never special-case this branch.
        n = min_ticks
        timestamps = np.arange(n + 1, dtype=np.float64) * dt
        s = np.zeros(n + 1, dtype=np.float64)
        s[-1] = 1.0
        s_dot = np.zeros(n + 1, dtype=np.float64)
        s_ddot = np.zeros(n + 1, dtype=np.float64)
        return ScalarProfile(timestamps, s, s_dot, s_ddot, float(timestamps[-1]), False, 0.0, 0.0)

    v_terms = []
    a_terms = []
    if not zero_length:
        v_terms.append(v_max / length_m)
        a_terms.append(cfg.a_max / length_m)
    if not zero_angle:
        v_terms.append(cfg.omega_max / angle_rad)
        a_terms.append(cfg.alpha_max / angle_rad)
    s_dot_max = min(v_terms)
    s_ddot_max = min(a_terms)

    is_triangular = (s_dot_max**2 / s_ddot_max) >= 1.0
    if is_triangular:
        t_acc = float(np.sqrt(1.0 / s_ddot_max))
        raw_duration = 2.0 * t_acc
    else:
        t_acc = s_dot_max / s_ddot_max
        raw_duration = 1.0 / s_dot_max + s_dot_max / s_ddot_max

    # Grid-snap: at least enough ticks to cover raw_duration, and at least
    # min_ticks (see the retiming use above). The "-1e-9" guards against
    # ceil() bumping a duration that is *exactly* on a tick boundary up by one
    # extra tick due to ordinary float noise in raw_duration's arithmetic.
    n = max(min_ticks, int(np.ceil(raw_duration * rate_hz - 1e-9)))
    duration = n / rate_hz

    # Stretch factor k <= 1: scaling s_dot_max/s_ddot_max down by k/k**2 and
    # re-deriving t_acc from them is algebraically identical to evaluating the
    # original (unstretched) profile at a time rescaled by k -- see the
    # module docstring. t_acc/duration is invariant under this rescaling
    # (both t_acc and duration scale by the same 1/k), which is what keeps the
    # accel/cruise/decel phase *proportions* identical, just slower.
    k = raw_duration / duration
    s_dot_max *= k
    s_ddot_max *= k * k
    t_acc /= k

    timestamps = np.arange(n + 1, dtype=np.float64) * dt
    s, s_dot, s_ddot = _sample_scurve(timestamps, duration, t_acc, s_dot_max, s_ddot_max)

    # Force exact boundary values. The phase-boundary arithmetic above is
    # exact in principle but not bit-exact in float64; downstream code relies
    # on s(0) == 0.0 and s(T) == 1.0 *exactly* (e.g. p(1) must equal the
    # commanded target position bit-for-bit, not target - 1e-15), and s_dot
    # must be exactly 0 at both ends since every segment starts and stops at
    # rest by construction (point-to-point, no blending between waypoints).
    s[0], s[-1] = 0.0, 1.0
    s_dot[0], s_dot[-1] = 0.0, 0.0

    return ScalarProfile(
        timestamps, s, s_dot, s_ddot, duration, is_triangular, s_dot_max, s_ddot_max
    )