screws 0.1.0__py3-none-any.whl

screws/__init__.py

@@ -0,0 +1,37 @@
+"""screws: screw-theory robotics, after Lynch and Park, *Modern Robotics* (MR).
+
+``import screws as sc`` gives every chapter function flat (``sc.exp6``, ``sc.fk_space``),
+MR's CamelCase names as aliases (``sc.FKinSpace is sc.fk_space``), the ``Robot`` class,
+the ``robots`` that ship, the ``urdf`` loader and the ``testing`` helpers. The CoppeliaSim
+bridge is ``screws.coppelia``, imported only on request.
+"""
+
+from . import kinematics, robots, se3, so3, testing, urdf
+from ._version import __version__
+from .aliases import *
+from .aliases import ALIASES
+from .kinematics import *
+from .kinematics import IKResult
+from .robot import MissingInertias, Robot
+from .se3 import *
+from .so3 import *
+
+__all__ = (
+    [
+        "__version__",
+        "ALIASES",
+        "IKResult",
+        "MissingInertias",
+        "Robot",
+        "kinematics",
+        "robots",
+        "se3",
+        "so3",
+        "testing",
+        "urdf",
+    ]
+    + list(so3.__all__)
+    + list(se3.__all__)
+    + list(kinematics.__all__)
+    + list(ALIASES)
+)

screws/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

screws/aliases.py

@@ -0,0 +1,94 @@
+"""The Modern Robotics library's CamelCase names, bound to their screws equivalents.
+
+Where the signatures match, the alias is the very same function object (``FKinSpace is
+fk_space``). Where screws changed the interface (the IK functions return an IKResult), the
+alias is a thin wrapper with MR's exact signature and return. No deprecation warnings: the
+names exist so a student can type what the book says.
+"""
+
+from __future__ import annotations
+
+from . import kinematics, se3, so3
+
+#: MR name -> screws primary name.
+ALIASES: dict[str, str] = {
+    # chapter 3
+    "NearZero": "near_zero",
+    "Normalize": "normalize",
+    "RotInv": "rot_inv",
+    "VecToso3": "vec_to_so3",
+    "so3ToVec": "so3_to_vec",
+    "AxisAng3": "axis_angle3",
+    "MatrixExp3": "exp3",
+    "MatrixLog3": "log3",
+    "RpToTrans": "rp_to_transform",
+    "TransToRp": "transform_to_rp",
+    "TransInv": "transform_inv",
+    "VecTose3": "vec_to_se3",
+    "se3ToVec": "se3_to_vec",
+    "Adjoint": "adjoint",
+    "ScrewToAxis": "screw_axis",
+    "AxisAng6": "axis_angle6",
+    "MatrixExp6": "exp6",
+    "MatrixLog6": "log6",
+    "ProjectToSO3": "project_so3",
+    "ProjectToSE3": "project_se3",
+    "DistanceToSO3": "distance_so3",
+    "DistanceToSE3": "distance_se3",
+    "TestIfSO3": "is_so3",
+    "TestIfSE3": "is_se3",
+    # chapter 4
+    "FKinBody": "fk_body",
+    "FKinSpace": "fk_space",
+    # chapter 5
+    "JacobianBody": "jacobian_body",
+    "JacobianSpace": "jacobian_space",
+    # chapter 6
+    "IKinBody": "ik_body",
+    "IKinSpace": "ik_space",
+}
+
+# Identity aliases: the same object under MR's name.
+NearZero = so3.near_zero
+Normalize = so3.normalize
+RotInv = so3.rot_inv
+VecToso3 = so3.vec_to_so3
+so3ToVec = so3.so3_to_vec
+AxisAng3 = so3.axis_angle3
+MatrixExp3 = so3.exp3
+MatrixLog3 = so3.log3
+RpToTrans = se3.rp_to_transform
+TransToRp = se3.transform_to_rp
+TransInv = se3.transform_inv
+VecTose3 = se3.vec_to_se3
+se3ToVec = se3.se3_to_vec
+Adjoint = se3.adjoint
+ScrewToAxis = se3.screw_axis
+AxisAng6 = se3.axis_angle6
+MatrixExp6 = se3.exp6
+MatrixLog6 = se3.log6
+ProjectToSO3 = so3.project_so3
+ProjectToSE3 = se3.project_se3
+DistanceToSO3 = so3.distance_so3
+DistanceToSE3 = se3.distance_se3
+TestIfSO3 = so3.is_so3
+TestIfSE3 = se3.is_se3
+FKinBody = kinematics.fk_body
+FKinSpace = kinematics.fk_space
+JacobianBody = kinematics.jacobian_body
+JacobianSpace = kinematics.jacobian_space
+
+
+def IKinBody(Blist, M, T, thetalist0, eomg, ev):
+    """MR's IKinBody: returns (thetalist, success). Wraps screws.kinematics.ik_body."""
+    r = kinematics.ik_body(M, Blist, T, thetalist0, tol_omega=eomg, tol_v=ev)
+    return r.theta, r.converged
+
+
+def IKinSpace(Slist, M, T, thetalist0, eomg, ev):
+    """MR's IKinSpace: returns (thetalist, success). Wraps screws.kinematics.ik_space."""
+    r = kinematics.ik_space(M, Slist, T, thetalist0, tol_omega=eomg, tol_v=ev)
+    return r.theta, r.converged
+
+
+__all__ = list(ALIASES) + ["ALIASES"]

screws/coppelia/__init__.py

@@ -0,0 +1,8 @@
+"""The CoppeliaSim bridge. Requires ``screws[coppelia]``; ``import screws`` never loads it."""
+
+from ._sim import SimulatorNotRunning, connect
+from .arm import Arm
+from .log import Log
+from .scene import Scene
+
+__all__ = ["Arm", "Log", "Scene", "SimulatorNotRunning", "connect"]

screws/coppelia/_sim.py

@@ -0,0 +1,119 @@
+"""The thin layer over CoppeliaSim's ZMQ remote API: connecting, and converting between
+CoppeliaSim's flat matrix, pose and inertia encodings and 4x4 / 3x3 numpy arrays.
+
+Everything here is testable without a simulator: the bridge receives a ``sim`` object and
+never imports the client except inside connect().
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ..se3 import rp_to_transform
+
+__all__ = [
+    "SimulatorNotRunning",
+    "connect",
+    "inertia9_to_matrix",
+    "matrix12_to_transform",
+    "pose7_to_transform",
+    "transform_to_matrix12",
+    "transform_to_pose7",
+]
+
+STUDENT_SENTENCE = "Start CoppeliaSim and open a scene, then run this again."
+
+
+class SimulatorNotRunning(RuntimeError):
+    """CoppeliaSim did not answer on the remote API port."""
+
+
+def _new_client(host: str, port: int):
+    from coppeliasim_zmqremoteapi_client import RemoteAPIClient
+
+    return RemoteAPIClient(host, port)
+
+
+def connect(host: str = "localhost", port: int = 23000, *, timeout_s: float = 5.0):
+    """Connect to CoppeliaSim's ZMQ remote API and return its ``sim`` object.
+
+    Raises SimulatorNotRunning, with the sentence a student needs, if nothing answers.
+    """
+    socket = None
+    try:
+        client = _new_client(host, port)
+        socket = getattr(client, "socket", None)
+        # The client's own `timeout` is a server-side setting sent with the first request;
+        # its recv() blocks forever. Put a real receive timeout on the socket for the probe.
+        if socket is not None:
+            socket.setsockopt(_zmq().RCVTIMEO, int(timeout_s * 1000))
+            socket.setsockopt(_zmq().LINGER, 0)
+        sim = client.require("sim")
+        sim.getSimulationTime()
+    except Exception as exc:  # zmq.Again, ConnectionRefusedError, ...: all mean "not running"
+        if socket is not None:
+            socket.close()
+        raise SimulatorNotRunning(
+            f"No CoppeliaSim answered at {host}:{port} ({type(exc).__name__}). {STUDENT_SENTENCE}"
+        ) from exc
+    if socket is not None:
+        socket.setsockopt(_zmq().RCVTIMEO, -1)
+    sim._screws_client = client  # keep the client alive as long as sim is
+    return sim
+
+
+def _zmq():
+    import zmq
+
+    return zmq
+
+
+def matrix12_to_transform(m) -> np.ndarray:
+    """CoppeliaSim's 12 floats (rows 0-2 of T, row-major) to a 4x4 transformation."""
+    T = np.eye(4)
+    T[:3, :] = np.asarray(m, dtype=float).reshape(3, 4)
+    return T
+
+
+def transform_to_matrix12(T) -> list[float]:
+    """A 4x4 transformation to CoppeliaSim's 12 floats (rows 0-2, row-major)."""
+    return [float(x) for x in np.asarray(T, dtype=float)[:3, :].reshape(-1)]
+
+
+def pose7_to_transform(p) -> np.ndarray:
+    """CoppeliaSim's pose (x, y, z, qx, qy, qz, qw) to a 4x4 transformation."""
+    p = np.asarray(p, dtype=float)
+    x, y, z, w = p[3:7] / np.linalg.norm(p[3:7])
+    R = np.array(
+        [
+            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
+            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
+            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
+        ]
+    )
+    return rp_to_transform(R, p[:3])
+
+
+def transform_to_pose7(T) -> list[float]:
+    """A 4x4 transformation to CoppeliaSim's pose (x, y, z, qx, qy, qz, qw)."""
+    T = np.asarray(T, dtype=float)
+    R = T[:3, :3]
+    tr = np.trace(R)
+    if tr > 0:
+        s = np.sqrt(tr + 1.0) * 2
+        w, x, y, z = 0.25 * s, (R[2, 1] - R[1, 2]) / s, (R[0, 2] - R[2, 0]) / s, (R[1, 0] - R[0, 1]) / s
+    elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
+        s = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2
+        w, x, y, z = (R[2, 1] - R[1, 2]) / s, 0.25 * s, (R[0, 1] + R[1, 0]) / s, (R[0, 2] + R[2, 0]) / s
+    elif R[1, 1] > R[2, 2]:
+        s = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2
+        w, x, y, z = (R[0, 2] - R[2, 0]) / s, (R[0, 1] + R[1, 0]) / s, 0.25 * s, (R[1, 2] + R[2, 1]) / s
+    else:
+        s = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2
+        w, x, y, z = (R[1, 0] - R[0, 1]) / s, (R[0, 2] + R[2, 0]) / s, (R[1, 2] + R[2, 1]) / s, 0.25 * s
+    return [float(v) for v in (*T[:3, 3], x, y, z, w)]
+
+
+def inertia9_to_matrix(i) -> np.ndarray:
+    """CoppeliaSim's 9 inertia floats (row-major) to a 3x3 matrix."""
+    return np.asarray(i, dtype=float).reshape(3, 3)

screws/coppelia/arm.py

@@ -0,0 +1,230 @@
+"""Arm: the joints under one scene tree, read and commanded through ``sim``."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ..robot import Robot
+from ..se3 import prismatic_axis, revolute_axis, transform_inv
+
+__all__ = ["Arm"]
+
+_MODES = ("position", "velocity", "torque")
+
+
+class Arm:
+    """A serial chain in the scene: its joints ordered base to tip, and its tip object.
+
+    Joints are every joint under ``path``, ordered by depth, or the ones named in ``joints``
+    (aliases or paths, in order) when a gripper or other tool adds joints of its own. The tip is the first dummy
+    under ``path`` whose alias is "tip" or "ee" (or contains "tip"), else the last joint.
+    """
+
+    def __init__(self, scene, path: str, joints=None):
+        self.scene = scene
+        self.sim = scene.sim
+        self.path = path
+        self.base = scene._handle(path)
+        self.alias = self.sim.getObjectAlias(self.base, -1)
+        found = list(self.sim.getObjectsInTree(self.base, self.sim.sceneobject_joint, 0))
+        found.sort(key=self._depth)
+        if joints is None:
+            joints = found
+        else:
+            by_alias = {self.sim.getObjectAlias(h, -1): h for h in found}
+            picked = []
+            for j in joints:
+                if isinstance(j, str) and j in by_alias:
+                    picked.append(by_alias[j])
+                elif isinstance(j, str):
+                    try:
+                        picked.append(scene._handle(j))
+                    except LookupError as exc:
+                        raise LookupError(
+                            f"no joint {j!r} under {path}; found {sorted(by_alias)}"
+                        ) from exc
+                else:
+                    picked.append(int(j))
+            joints = picked
+        self.handles: tuple[int, ...] = tuple(joints)
+        self.joint_names: tuple[str, ...] = tuple(self.sim.getObjectAlias(h, -1) for h in joints)
+        self.tip, self.tip_alias = self._find_tip()
+        self._mode: str | None = None
+
+    def _depth(self, h: int) -> int:
+        d = 0
+        while h != self.sim.handle_world and h != self.base:
+            h = self.sim.getObjectParent(h)
+            d += 1
+        return d
+
+    def _find_tip(self) -> tuple[int, str]:
+        dummies = list(self.sim.getObjectsInTree(self.base, self.sim.sceneobject_dummy, 0))
+        named = {self.sim.getObjectAlias(h, -1): h for h in dummies}
+        for want in ("tip", "ee", "connection"):
+            if want in named:
+                return named[want], want
+        for alias, h in named.items():
+            if any(key in alias.lower() for key in ("tip", "connection")):
+                return h, alias
+        if not self.handles:
+            raise LookupError(f"{self.path} has no joints and no tip dummy")
+        # Fall back to the last joint's first non-joint child (the last link), else the joint.
+        last = self.handles[-1]
+        for kind in (self.sim.sceneobject_shape, self.sim.sceneobject_dummy):
+            for h in self.sim.getObjectsInTree(last, kind, 1 + 2):  # exclude base, first children only
+                return h, self.sim.getObjectAlias(h, -1)
+        return last, self.joint_names[-1]
+
+    @property
+    def n(self) -> int:
+        return len(self.handles)
+
+    # ----- reading ----------------------------------------------------------------------
+
+    def theta(self) -> np.ndarray:
+        """Joint positions (rad or m)."""
+        return np.array([self.sim.getJointPosition(h) for h in self.handles], dtype=float)
+
+    def dtheta(self) -> np.ndarray:
+        """Joint velocities."""
+        return np.array([self.sim.getJointVelocity(h) for h in self.handles], dtype=float)
+
+    def tau(self) -> np.ndarray:
+        """Measured joint forces or torques; NaN for a joint that reports none (not dynamic)."""
+        out = []
+        for h in self.handles:
+            try:
+                out.append(float(self.sim.getJointForce(h)))
+            except Exception:  # noqa: BLE001 - the remote API raises a plain Exception
+                out.append(float("nan"))
+        return np.array(out, dtype=float)
+
+    def tip_frame(self) -> np.ndarray:
+        """The tip's 4x4 configuration in the world frame."""
+        return self.scene.frame(self.tip)
+
+    def joint_frames(self) -> list[np.ndarray]:
+        """Each joint's 4x4 frame in the world frame; the joint axis is the local z axis."""
+        return [self.scene.frame(h) for h in self.handles]
+
+    def joint_types(self) -> tuple[str, ...]:
+        return tuple(
+            "prismatic" if self.sim.getJointType(h) == self.sim.joint_prismatic else "revolute"
+            for h in self.handles
+        )
+
+    def joint_limits(self) -> np.ndarray | None:
+        """nx2 limits from the joints' intervals, or None if any joint is cyclic."""
+        out = []
+        for h in self.handles:
+            cyclic, interval = self.sim.getJointInterval(h)
+            if cyclic:
+                return None
+            out.append([interval[0], interval[0] + interval[1]])
+        return np.array(out, dtype=float)
+
+    # ----- commanding -------------------------------------------------------------------
+
+    def mode(self, kind: str) -> None:
+        """Set every joint's dynamic control mode: "position", "velocity" or "torque"."""
+        if kind not in _MODES:
+            raise ValueError(f"mode must be one of {_MODES}; got {kind!r}")
+        value = {
+            "position": self.sim.jointdynctrl_position,
+            "velocity": self.sim.jointdynctrl_velocity,
+            "torque": self.sim.jointdynctrl_force,
+        }[kind]
+        for h in self.handles:
+            self.sim.setObjectInt32Param(h, self.sim.jointintparam_dynctrlmode, value)
+        self._mode = kind
+
+    def _require(self, kind: str) -> None:
+        if self._mode != kind:
+            raise RuntimeError(
+                f"the arm is in {self._mode!r} mode; call arm.mode({kind!r}) before a {kind} command"
+            )
+
+    def command(self, u) -> None:
+        """Dispatch on the current mode: positions, velocities or torques."""
+        if self._mode is None:
+            raise RuntimeError('set a mode first: arm.mode("position" | "velocity" | "torque")')
+        dispatch = {
+            "position": self.command_positions,
+            "velocity": self.command_velocities,
+            "torque": self.command_torques,
+        }
+        dispatch[self._mode](u)
+
+    def _vector(self, u) -> np.ndarray:
+        u = np.asarray(u, dtype=float).reshape(-1)
+        if u.shape[0] != self.n:
+            raise ValueError(f"got {u.shape[0]} values for {self.n} joints")
+        return u
+
+    def command_positions(self, theta) -> None:
+        self._require("position")
+        for h, v in zip(self.handles, self._vector(theta)):
+            self.sim.setJointTargetPosition(h, float(v))
+
+    def command_velocities(self, dtheta) -> None:
+        self._require("velocity")
+        for h, v in zip(self.handles, self._vector(dtheta)):
+            self.sim.setJointTargetVelocity(h, float(v))
+
+    def command_torques(self, tau) -> None:
+        """Force mode: the signed joint force or torque is applied directly
+        (sim.setJointTargetForce with signedValue, CoppeliaSim 4.3+)."""
+        self._require("torque")
+        for h, v in zip(self.handles, self._vector(tau)):
+            self.sim.setJointTargetForce(h, float(v), True)
+
+    def teleport(self, theta) -> None:
+        """Set joint positions directly, without physics: for animating IK iterates."""
+        for h, v in zip(self.handles, self._vector(theta)):
+            self.sim.setJointPosition(h, float(v))
+
+    # ----- the robot off the scene ------------------------------------------------------
+
+    def robot(self, *, inertias: bool = False, relative_to: str = "world") -> Robot:
+        """A screws.Robot read from the scene at the zero position.
+
+        M is the tip frame; joint i's screw axis has omega = its frame's z axis and
+        q = its origin (notes 4.1: v = -omega x q). relative_to="world" (default) makes {s}
+        CoppeliaSim's world frame; relative_to="base" makes {s} the frame of the object at
+        ``path``, so the result is independent of where the model stands in the scene.
+        The arm is teleported to zero for the reading and put back afterwards, so call this
+        before starting the simulation or accept a jump. inertias=True arrives in 0.2.
+        """
+        if inertias:
+            raise NotImplementedError("scene inertias arrive in screws 0.2")
+        if relative_to not in ("world", "base"):
+            raise ValueError(f'relative_to must be "world" or "base"; got {relative_to!r}')
+        here = self.theta()
+        self.teleport(np.zeros(self.n))
+        try:
+            frames = self.joint_frames()
+            M = self.tip_frame()
+            if relative_to == "base":
+                T_ws_inv = transform_inv(self.scene.frame(self.base))
+                frames = [T_ws_inv @ F for F in frames]
+                M = T_ws_inv @ M
+        finally:
+            self.teleport(here)
+        axes = []
+        for F, kind in zip(frames, self.joint_types()):
+            z, q = F[:3, 2], F[:3, 3]
+            axes.append(prismatic_axis(z) if kind == "prismatic" else revolute_axis(q, z))
+        return Robot(
+            name=self.alias,
+            M=M,
+            S=np.column_stack(axes),
+            joint_types=self.joint_types(),
+            joint_names=self.joint_names,
+            joint_limits=self.joint_limits(),
+            joint_frames_home=tuple(frames),
+        )
+
+    def __repr__(self) -> str:
+        return f"Arm({self.path!r}, joints={list(self.joint_names)}, tip={self.tip_alias!r})"
+

screws/coppelia/log.py

@@ -0,0 +1,105 @@
+"""Log: what a Scene records at every step."""
+
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import numpy as np
+
+__all__ = ["Log"]
+
+
+class Log:
+    """Per-step records of simulated time, joint state, command and end-effector frame.
+
+    Arrays: t (N,), theta (N, n), dtheta (N, n), tau (N, n), command (N, n) and T_sb
+    (N, 4, 4). Empty arrays before the first record.
+    """
+
+    def __init__(self):
+        self._t: list[float] = []
+        self._theta: list[np.ndarray] = []
+        self._dtheta: list[np.ndarray] = []
+        self._tau: list[np.ndarray] = []
+        self._command: list[np.ndarray] = []
+        self._T_sb: list[np.ndarray] = []
+
+    def record(self, t, theta, dtheta, tau, command, T_sb) -> None:
+        self._t.append(float(t))
+        self._theta.append(np.asarray(theta, dtype=float))
+        self._dtheta.append(np.asarray(dtheta, dtype=float))
+        self._tau.append(np.asarray(tau, dtype=float))
+        n = len(self._theta[-1])
+        self._command.append(
+            np.full(n, np.nan) if command is None else np.asarray(command, dtype=float)
+        )
+        self._T_sb.append(np.asarray(T_sb, dtype=float))
+
+    def __len__(self) -> int:
+        return len(self._t)
+
+    @property
+    def t(self) -> np.ndarray:
+        return np.array(self._t)
+
+    @property
+    def theta(self) -> np.ndarray:
+        return np.array(self._theta)
+
+    @property
+    def dtheta(self) -> np.ndarray:
+        return np.array(self._dtheta)
+
+    @property
+    def tau(self) -> np.ndarray:
+        return np.array(self._tau)
+
+    @property
+    def command(self) -> np.ndarray:
+        return np.array(self._command)
+
+    @property
+    def T_sb(self) -> np.ndarray:
+        return np.array(self._T_sb)
+
+    def to_csv(self, path) -> None:
+        """Write t, theta_i, dtheta_i, tau_i, command_i and the 12 entries of T_sb per row."""
+        n = self.theta.shape[1] if len(self) else 0
+        header = (
+            ["t"]
+            + [f"theta_{i + 1}" for i in range(n)]
+            + [f"dtheta_{i + 1}" for i in range(n)]
+            + [f"tau_{i + 1}" for i in range(n)]
+            + [f"command_{i + 1}" for i in range(n)]
+            + [f"T_{r}{c}" for r in range(3) for c in range(4)]
+        )
+        with Path(path).open("w", newline="") as f:
+            w = csv.writer(f)
+            w.writerow(header)
+            for k in range(len(self)):
+                w.writerow(
+                    [self._t[k], *self._theta[k], *self._dtheta[k], *self._tau[k],
+                     *self._command[k], *self._T_sb[k][:3, :].reshape(-1)]
+                )
+
+    def to_mr_csv(self, path) -> None:
+        """Write joint angles only, one row per step, no header: the MR wiki scenes' format."""
+        np.savetxt(path, self.theta, delimiter=",")
+
+    def plot(self):
+        """One axis per quantity (theta, dtheta, tau, command). Needs matplotlib."""
+        import matplotlib.pyplot as plt
+
+        fig, axes = plt.subplots(4, 1, sharex=True, figsize=(8, 9))
+        for ax, (name, data) in zip(
+            axes, [("theta", self.theta), ("dtheta", self.dtheta), ("tau", self.tau),
+                   ("command", self.command)]
+        ):
+            if len(self):
+                ax.plot(self.t, data)
+            ax.set_ylabel(name)
+            ax.grid(True)
+        axes[-1].set_xlabel("t (s)")
+        fig.tight_layout()
+        return fig