cmvr-head-client/s_curve.py

import importlib.util
from pathlib import Path
import sys


REPO_ROOT = Path(__file__).resolve().parent


class SCurveError(RuntimeError):
    pass


def _default_module_candidates():
    search_roots = [
        REPO_ROOT / "cpp" / "common" / "curve" / "build",
        REPO_ROOT / "cpp" / "common" / "curve" / "build" / "Release",
    ]
    patterns = [
        "_s_curve_native*.pyd",
        "_s_curve_native*.so",
    ]
    candidates = []
    for root in search_roots:
        if not root.exists():
            continue
        for pattern in patterns:
            candidates.extend(sorted(root.glob(pattern)))
    return candidates


def _resolve_module_path(module_path=None):
    if module_path:
        path = Path(module_path)
        if path.exists():
            return path
        raise FileNotFoundError(f"SCurve native module not found: {path}")

    for candidate in _default_module_candidates():
        if candidate.exists():
            return candidate

    candidate_text = "\n".join(str(path) for path in _default_module_candidates())
    raise FileNotFoundError(
        "SCurve native module not found. Build it first with:\n"
        "powershell -ExecutionPolicy Bypass -File .\\cpp\\common\\curve\\build_python.ps1 "
        "-PythonExecutable 'C:\\Users\\Administrator\\miniconda3\\envs\\cmvr-head-client\\python.exe'\n"
        f"Expected one of:\n{candidate_text}"
    )


class SCurve:
    def __init__(self, max_velocity=3.0, max_acceleration=10.0, max_jerk=50.0, module_path=None, dll_path=None):
        resolved_path = _resolve_module_path(module_path or dll_path)
        try:
            self._module = self._load_native_module(resolved_path)
            self._native = self._module.SCurve(max_velocity, max_acceleration, max_jerk)
        except Exception as exc:
            raise SCurveError(f"Failed to create pybind11 SCurve instance: {exc}") from exc

    @staticmethod
    def _load_native_module(module_path):
        module_name = "_s_curve_native"
        if module_name in sys.modules:
            return sys.modules[module_name]
        spec = importlib.util.spec_from_file_location(module_name, module_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Unable to load module from {module_path}")
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
        return module

    @staticmethod
    def _wrap_native_call(action, fn, *args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except Exception as exc:
            raise SCurveError(f"{action} failed: {exc}") from exc

    def setConstraints(self, max_velocity, max_acceleration, max_jerk):
        return self._wrap_native_call(
            "setConstraints",
            self._native.setConstraints,
            max_velocity,
            max_acceleration,
            max_jerk,
        )

    def getConstraints(self):
        return self._wrap_native_call(
            "getConstraints",
            self._native.getConstraints,
        )

    def calculateProfile(self, start_position, end_position, start_velocity=0.0, end_velocity=0.0):
        return self._wrap_native_call(
            "calculateProfile",
            self._native.calculateProfile,
            start_position,
            end_position,
            start_velocity,
            end_velocity,
        )

    def getPositionAtTime(self, profile, t):
        return self._wrap_native_call("getPositionAtTime", self._native.getPositionAtTime, profile, t)

    def getVelocityAtTime(self, profile, t):
        return self._wrap_native_call("getVelocityAtTime", self._native.getVelocityAtTime, profile, t)

    def getAccelerationAtTime(self, profile, t):
        return self._wrap_native_call(
            "getAccelerationAtTime",
            self._native.getAccelerationAtTime,
            profile,
            t,
        )

    def getJerkAtTime(self, profile, t):
        return self._wrap_native_call("getJerkAtTime", self._native.getJerkAtTime, profile, t)

    def sampleTrajectory(self, profile, dt):
        return self._wrap_native_call(
            "sampleTrajectory",
            self._native.sampleTrajectory,
            profile,
            dt,
        )


class SCurvePositionPlanner1D:
    def __init__(self, max_velocity=3.0, max_acceleration=10.0, max_jerk=50.0, module_path=None, dll_path=None):
        resolved_path = _resolve_module_path(module_path or dll_path)
        try:
            self._module = SCurve._load_native_module(resolved_path)
            self._native = self._module.SCurvePositionPlanner1D(
                max_velocity,
                max_acceleration,
                max_jerk,
            )
        except Exception as exc:
            raise SCurveError(f"Failed to create pybind11 SCurvePositionPlanner1D instance: {exc}") from exc

    def setConstraints(self, max_velocity, max_acceleration, max_jerk):
        return SCurve._wrap_native_call(
            "setConstraints",
            self._native.setConstraints,
            max_velocity,
            max_acceleration,
            max_jerk,
        )

    def getConstraints(self):
        return SCurve._wrap_native_call("getConstraints", self._native.getConstraints)

    def setPositionGain(self, position_gain):
        return SCurve._wrap_native_call("setPositionGain", self._native.setPositionGain, position_gain)

    def getPositionGain(self):
        return SCurve._wrap_native_call("getPositionGain", self._native.getPositionGain)

    def initialize(self, position, velocity=0.0, acceleration=0.0):
        return SCurve._wrap_native_call(
            "initialize",
            self._native.initialize,
            position,
            velocity,
            acceleration,
        )

    def reset(self):
        return SCurve._wrap_native_call("reset", self._native.reset)

    def setTarget(self, target_position):
        return SCurve._wrap_native_call("setTarget", self._native.setTarget, target_position)

    def update(self, dt):
        return SCurve._wrap_native_call("update", self._native.update, dt)

    def getState(self):
        return SCurve._wrap_native_call("getState", self._native.getState)

    def sampleTrajectory(
        self,
        target_position,
        dt,
        duration,
        initial_position=0.0,
        initial_velocity=0.0,
        initial_acceleration=0.0,
        position_gain=None,
    ):
        if dt <= 0:
            raise ValueError("dt must be > 0")
        if duration < 0:
            raise ValueError("duration must be >= 0")

        self.initialize(initial_position, initial_velocity, initial_acceleration)
        if position_gain is not None:
            self.setPositionGain(position_gain)
        self.setTarget(target_position)

        times = [0.0]
        state = self.getState()
        positions = [state.position]
        velocities = [state.velocity]
        accelerations = [state.acceleration]
        jerks = [state.jerk]
        targets = [state.target_position]
        moving = [state.is_moving]

        t = 0.0
        while t + dt <= duration + 1e-12:
            t = round(t + dt, 12)
            self.update(dt)
            state = self.getState()
            times.append(t)
            positions.append(state.position)
            velocities.append(state.velocity)
            accelerations.append(state.acceleration)
            jerks.append(state.jerk)
            targets.append(state.target_position)
            moving.append(state.is_moving)

        if abs(times[-1] - duration) > 1e-9:
            final_dt = duration - times[-1]
            if final_dt > 1e-12:
                self.update(final_dt)
                state = self.getState()
                times.append(duration)
                positions.append(state.position)
                velocities.append(state.velocity)
                accelerations.append(state.acceleration)
                jerks.append(state.jerk)
                targets.append(state.target_position)
                moving.append(state.is_moving)

        return {
            "times": times,
            "positions": positions,
            "velocities": velocities,
            "accelerations": accelerations,
            "jerks": jerks,
            "targets": targets,
            "moving": moving,
        }