virtualmodelcontrol 0.1.0__py3-none-any.whl

virtualmodelcontrol/mechanisms/components/inertance.py

@@ -0,0 +1,56 @@
+"""Inertance components: masses and inertias, with kinetic energy ½ ẏᵀ M ẏ."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from ...core.registry import register
+from ...core.units import KG, inertance_unit
+from ..coordinates.base import Context, Coordinate
+from .base import Component
+
+
+@register("component", "point_mass")
+class PointMass(Component):
+    """A mass m [kg] at a point coordinate; a ``design`` Param (a physical property)."""
+
+    kind = "inertance"
+
+    def __init__(self, coord: Coordinate, mass: Any) -> None:
+        super().__init__(coord)
+        self.mass = self._param("mass", mass, unit=KG, scope="design")
+
+    def inertance(self, ctx: Context, y: Any) -> Any:
+        """m I."""
+        return ctx.param(self.mass) * ca.DM.eye(self.coord.dim)
+
+
+@register("component", "inertance")
+class Inertance(Component):
+    """Constant inertance M on a coordinate: a scalar, one value per axis, or a matrix.
+
+    An ``episode`` Param: a virtual flywheel's inertia, say [kg·m² on angles].
+    """
+
+    kind = "inertance"
+
+    def __init__(self, coord: Coordinate, inertance: Any) -> None:
+        super().__init__(coord)
+        matrix = np.ndim(getattr(inertance, "value", inertance)) == 2
+        self.inertia = self._param(
+            "inertance",
+            inertance,
+            unit=inertance_unit(coord.unit),
+            scope="episode",
+            bounds=(-np.inf, np.inf) if matrix else (0.0, np.inf),
+        )
+
+    def inertance(self, ctx: Context, y: Any) -> Any:
+        """M as a (dim, dim) matrix."""
+        M = ctx.param(self.inertia)
+        if M.shape[1] == 1:
+            return ca.diag(M * ca.DM.ones(self.coord.dim, 1))
+        return M

virtualmodelcontrol/mechanisms/components/sources.py

@@ -0,0 +1,63 @@
+"""Source components: forces whose power is metered in the energy accounting."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from ...core.registry import register
+from ...core.units import M_S2, force_unit
+from ..coordinates.base import Context, Coordinate
+from ..coordinates.ops import Stack
+from .base import Component
+from .inertance import PointMass
+
+
+@register("component", "force_source")
+class ForceSource(Component):
+    """A force f on a coordinate, held in a live Param; its power f·ẏ is metered."""
+
+    kind = "source"
+
+    def __init__(self, coord: Coordinate, force: Any) -> None:
+        super().__init__(coord)
+        free = (-np.inf, np.inf)
+        self.force_value = self._param(
+            "force", force, unit=force_unit(coord.unit), scope="stage", bounds=free
+        )
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """The Param, as a column."""
+        return ca.reshape(ctx.param(self.force_value), self.coord.dim, 1)
+
+
+@register("component", "gravity_compensation")
+class GravityCompensation(Component):
+    """Forces cancelling gravity on a robot's point masses, f_i = −m_i g (a metered source).
+
+    ``gravity`` g is a vector in the robot's base frame [m/s²]; by default the robot's own
+    ``gravity`` Param, shared so both always agree.
+    """
+
+    kind = "source"
+
+    def __init__(self, robot: Any, gravity: Any = None) -> None:
+        masses = [c for c in robot.components.values() if isinstance(c, PointMass)]
+        if not masses:
+            raise ValueError(f"robot {robot.name!r} has no PointMass components")
+        super().__init__(Stack(*[m.coord for m in masses]))
+        self.masses = masses
+        if gravity is None:
+            if "gravity" not in robot.params:
+                raise ValueError(f"robot {robot.name!r} has no 'gravity' Param; pass gravity=...")
+            gravity = robot.params["gravity"]
+        self.gravity = self._param(
+            "gravity", gravity, unit=M_S2, scope="design", bounds=(-np.inf, np.inf)
+        )
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """−m_i g on each mass, stacked."""
+        g = ca.reshape(ctx.param(self.gravity), 3, 1)
+        return ca.vertcat(*[-ctx.param(m.mass) * g for m in self.masses])

virtualmodelcontrol/mechanisms/components/storage.py

@@ -0,0 +1,272 @@
+"""Storage components: springs with an energy V(y) ≥ 0 and the force f = −∂V/∂y.
+
+The deflection is y = x − x_ref, so a spring pulls x towards x_ref. Gains are Params: ``stage``
+scope (live) for stiffnesses and shapes, ``episode`` for limits and exponents.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from ...core.registry import register
+from ...core.symbolic import logcosh, quad, smooth_norm
+from ...core.units import stiffness_unit
+from ..coordinates.base import Context, Coordinate
+from .base import Component, scaled
+
+
+@register("component", "linear_spring")
+class LinearSpring(Component):
+    """Linear spring, V = ½ yᵀ K y and f = −K y.
+
+    ``stiffness`` K is a scalar, one value per axis, or a matrix (its symmetric part acts)
+    [N/m, or N·m/rad on angles].
+    """
+
+    kind = "storage"
+
+    def __init__(self, coord: Coordinate, stiffness: Any) -> None:
+        super().__init__(coord)
+        matrix = np.ndim(getattr(stiffness, "value", stiffness)) == 2
+        self.stiffness = self._param(
+            "stiffness",
+            stiffness,
+            unit=stiffness_unit(coord.unit),
+            scope="stage",
+            bounds=(-np.inf, np.inf) if matrix else (0.0, np.inf),
+        )
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """½ yᵀ K y."""
+        return 0.5 * ca.dot(y, scaled(ctx.param(self.stiffness), y))
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """−K y (symmetric part of K for a matrix)."""
+        K = ctx.param(self.stiffness)
+        if K.shape[1] == 1:
+            return -K * y
+        return -0.5 * (ca.mtimes(K, y) + ca.mtimes(K.T, y))
+
+
+@register("component", "tanh_spring")
+class TanhSpring(Component):
+    """Saturating spring, per axis: f = −F tanh(k y / F), within ±F on each axis.
+
+    ``stiffness`` k is the slope at y = 0, a scalar or one value per axis; ``max_force`` F [N].
+    """
+
+    kind = "storage"
+
+    def __init__(self, coord: Coordinate, stiffness: Any, max_force: Any) -> None:
+        super().__init__(coord)
+        self.stiffness = self._param(
+            "stiffness", stiffness, unit=stiffness_unit(coord.unit), scope="stage"
+        )
+        self.max_force = self._param("max_force", max_force, unit="N", scope="stage")
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """Σ (F²/k) log cosh(k y / F); zero on axes with k = 0."""
+        k = ctx.param(self.stiffness) * ca.DM.ones(self.coord.dim, 1)
+        F = ctx.param(self.max_force)
+        zero = ca.DM.zeros(self.coord.dim, 1)
+        return ca.sum1(ca.if_else(k > 0, F**2 / k * logcosh(k * y / F), zero))
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """−F tanh(k y / F)."""
+        k, F = ctx.param(self.stiffness), ctx.param(self.max_force)
+        return -F * ca.tanh(k * y / F)
+
+
+@register("component", "gaussian_spring")
+class GaussianSpring(Component):
+    """Repulsive Gaussian, V = A σ² exp(−‖y‖²/2σ²) and f = A exp(−‖y‖²/2σ²) y.
+
+    With y = x − x_obstacle it pushes x away from the obstacle. ``strength`` A [N/m], ``sigma``
+    σ in the coordinate's unit.
+    """
+
+    kind = "storage"
+
+    def __init__(self, coord: Coordinate, strength: Any, sigma: Any) -> None:
+        super().__init__(coord)
+        self.strength = self._param(
+            "strength", strength, unit=stiffness_unit(coord.unit), scope="stage"
+        )
+        self.sigma = self._param("sigma", sigma, unit=coord.unit, scope="stage")
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """A σ² exp(−‖y‖²/2σ²)."""
+        A, s = ctx.param(self.strength), ctx.param(self.sigma)
+        return A * s**2 * ca.exp(-ca.sumsqr(y) / (2 * s**2))
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """A exp(−‖y‖²/2σ²) y."""
+        A, s = ctx.param(self.strength), ctx.param(self.sigma)
+        return A * ca.exp(-ca.sumsqr(y) / (2 * s**2)) * y
+
+
+@register("component", "sigmoid_spring")
+class SigmoidSpring(Component):
+    """Spring stiffening with distance: f = −k(d) y, k(d) = k_min + (k_max − k_min) σ(α(d − d₀)).
+
+    d is each axis' magnitude (``element_wise``) or the norm ‖y‖. The energy has no closed form
+    and is integrated by Gauss–Legendre quadrature with ``nodes`` points; the force is exact.
+    """
+
+    kind = "storage"
+
+    def __init__(
+        self,
+        coord: Coordinate,
+        k_min: Any,
+        k_max: Any,
+        threshold: Any,
+        alpha: Any,
+        *,
+        element_wise: bool = True,
+        nodes: int = 32,
+    ) -> None:
+        super().__init__(coord)
+        unit = stiffness_unit(coord.unit)
+        self.k_min = self._param("k_min", k_min, unit=unit, scope="stage")
+        self.k_max = self._param("k_max", k_max, unit=unit, scope="stage")
+        self.threshold = self._param("threshold", threshold, unit=coord.unit, scope="stage")
+        inv = f"1/{coord.unit}" if coord.unit else ""
+        self.alpha = self._param("alpha", alpha, unit=inv, scope="stage")
+        self.element_wise = element_wise
+        self.nodes = nodes
+
+    def _stiffness(self, ctx: Context, d: Any) -> Any:
+        k0, k1 = ctx.param(self.k_min), ctx.param(self.k_max)
+        d0, a = ctx.param(self.threshold), ctx.param(self.alpha)
+        return k0 + (k1 - k0) / (1 + ca.exp(-a * (d - d0)))
+
+    def _distance(self, y: Any) -> Any:
+        return ca.fabs(y) if self.element_wise else smooth_norm(y)
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """Σ ∫₀^d k(r) r dr, by quadrature."""
+        d = self._distance(y)
+        integral = quad(lambda t: self._stiffness(ctx, d * t) * t, 0.0, 1.0, self.nodes)
+        return ca.sum1(d**2 * integral)
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """−k(d) y."""
+        return -self._stiffness(ctx, self._distance(y)) * y
+
+
+@register("component", "polynomial_spring")
+class PolynomialSpring(Component):
+    """Spring with stiffness k = K (d / d₀)^n: f = −k y, V = K d^(n+2) / ((n+2) d₀ⁿ).
+
+    d is each axis' magnitude (``element_wise``) or the norm ‖y‖. ``stiffness`` K is the
+    stiffness at d = d₀.
+    """
+
+    kind = "storage"
+
+    def __init__(
+        self,
+        coord: Coordinate,
+        stiffness: Any,
+        order: Any,
+        dist_norm: Any,
+        *,
+        element_wise: bool = True,
+    ) -> None:
+        super().__init__(coord)
+        self.stiffness = self._param(
+            "stiffness", stiffness, unit=stiffness_unit(coord.unit), scope="stage"
+        )
+        self.order = self._param("order", order, unit="", scope="episode")
+        self.dist_norm = self._param("dist_norm", dist_norm, unit=coord.unit, scope="stage")
+        self.element_wise = element_wise
+
+    def _distance(self, y: Any) -> Any:
+        return ca.fabs(y) if self.element_wise else smooth_norm(y)
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """Σ K d^(n+2) / ((n+2) d₀ⁿ)."""
+        K, n, d0 = ctx.param(self.stiffness), ctx.param(self.order), ctx.param(self.dist_norm)
+        return ca.sum1(K * self._distance(y) ** (n + 2) / ((n + 2) * d0**n))
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """−K (d / d₀)ⁿ y."""
+        K, n, d0 = ctx.param(self.stiffness), ctx.param(self.order), ctx.param(self.dist_norm)
+        return -K * (self._distance(y) / d0) ** n * y
+
+
+@register("component", "limit_spring")
+class LimitSpring(Component):
+    """Zero force inside [lower, upper] and linear outside, per axis (a joint-limit spring).
+
+    V = ½ k (max(0, lower − y)² + max(0, y − upper)²).
+    """
+
+    kind = "storage"
+
+    def __init__(self, coord: Coordinate, stiffness: Any, lower: Any, upper: Any) -> None:
+        super().__init__(coord)
+        self.stiffness = self._param(
+            "stiffness", stiffness, unit=stiffness_unit(coord.unit), scope="stage"
+        )
+        free = (-np.inf, np.inf)
+        self.lower = self._param("lower", lower, unit=coord.unit, scope="episode", bounds=free)
+        self.upper = self._param("upper", upper, unit=coord.unit, scope="episode", bounds=free)
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """½ k (max(0, lower − y)² + max(0, y − upper)²)."""
+        k = ctx.param(self.stiffness)
+        below = ca.fmax(0, ctx.param(self.lower) - y)
+        above = ca.fmax(0, y - ctx.param(self.upper))
+        return ca.sum1(0.5 * k * (below**2 + above**2))
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """k (max(0, lower − y) − max(0, y − upper))."""
+        k = ctx.param(self.stiffness)
+        below = ca.fmax(0, ctx.param(self.lower) - y)
+        above = ca.fmax(0, y - ctx.param(self.upper))
+        return k * (below - above)
+
+
+@register("component", "gravity")
+class Gravity(Component):
+    """Gravity on a robot's point masses: V = −Σ m_i gᵀ p_i, f_i = m_i g (a physical component).
+
+    Unlike a spring, V has no lower bound. ``gravity`` g [m/s², base frame] defaults to the
+    robot's ``gravity`` Param.
+    """
+
+    kind = "storage"
+
+    def __init__(self, robot: Any, gravity: Any = None) -> None:
+        from ..coordinates.ops import Stack
+        from .inertance import PointMass
+
+        masses = [c for c in robot.components.values() if isinstance(c, PointMass)]
+        if not masses:
+            raise ValueError(f"robot {robot.name!r} has no PointMass components")
+        super().__init__(Stack(*[m.coord for m in masses]))
+        self.masses = masses
+        if gravity is None:
+            if "gravity" not in robot.params:
+                raise ValueError(f"robot {robot.name!r} has no 'gravity' Param; pass gravity=...")
+            gravity = robot.params["gravity"]
+        self.gravity = self._param(
+            "gravity", gravity, unit="m/s^2", scope="design", bounds=(-np.inf, np.inf)
+        )
+
+    def _weights(self, ctx: Context) -> Any:
+        g = ca.reshape(ctx.param(self.gravity), 3, 1)
+        return ca.vertcat(*[ctx.param(m.mass) * g for m in self.masses])
+
+    def energy(self, ctx: Context, y: Any) -> Any:
+        """−Σ m_i gᵀ p_i."""
+        return -ca.dot(self._weights(ctx), y)
+
+    def force(self, ctx: Context, y: Any, yd: Any) -> Any:
+        """m_i g on each mass, stacked."""
+        return self._weights(ctx)

virtualmodelcontrol/mechanisms/coordinates/__init__.py

@@ -0,0 +1,24 @@
+"""Coordinates: what components act on, built from q, the virtual states z, Params and time."""
+
+from .base import Context, Coordinate, as_coordinate, walk
+from .frames import FramePoint
+from .joints import Joint, State
+from .ops import Custom, Difference, Norm, Projection, Slice, Stack
+from .references import Ref
+
+__all__ = [
+    "Context",
+    "Coordinate",
+    "Custom",
+    "Difference",
+    "FramePoint",
+    "Joint",
+    "Norm",
+    "Projection",
+    "Ref",
+    "Slice",
+    "Stack",
+    "State",
+    "as_coordinate",
+    "walk",
+]

virtualmodelcontrol/mechanisms/coordinates/base.py

@@ -0,0 +1,117 @@
+"""Coordinates: quantities y(q, z, p, t) that components act on, and where they are evaluated."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping
+from typing import Any
+
+from ...core.params import Binding, Param
+
+
+class Context:
+    """Symbols a coordinate is evaluated on: robot configuration q, virtual positions z, time t.
+
+    Params resolve through ``binding`` (live symbols or folded constants). Values are memoized
+    per coordinate object, so a coordinate shared by several components is built once.
+    """
+
+    def __init__(
+        self,
+        q: Any,
+        binding: Binding,
+        z: Any = None,
+        t: Any = None,
+        states: Mapping[int, slice] | None = None,
+    ) -> None:
+        self.q = q
+        self.z = z
+        self.t = t
+        self.binding = binding
+        self._states = dict(states or {})
+        self._cache: dict[int, Any] = {}
+
+    def param(self, param: Param) -> Any:
+        """Expression of a Param."""
+        return self.binding(param)
+
+    def view(self, params: Mapping[str, Param]) -> dict[str, Any]:
+        """Expressions of a model's Params, keyed by the model's own names."""
+        return self.binding.view(params)
+
+    def state(self, state: Coordinate) -> Any:
+        """Slice of z holding a virtual state."""
+        try:
+            return self.z[self._states[id(state)]]
+        except KeyError:
+            raise KeyError(f"{state!r} is not a state of this system") from None
+
+    def value(self, coord: Coordinate) -> Any:
+        """Value of ``coord`` (dim × 1), built once per context."""
+        key = id(coord)
+        if key not in self._cache:
+            self._cache[key] = coord.value(self)
+        return self._cache[key]
+
+
+class Coordinate:
+    """A quantity y with ``dim`` entries in ``unit``, computed from q, z, the Params and t.
+
+    Subtracting coordinates gives a ``Difference``; a plain array on either side becomes a
+    ``Ref`` (a live reference). Indexing gives a ``Slice``.
+    """
+
+    def __init__(self, dim: int, unit: str = "") -> None:
+        self.dim = int(dim)
+        self.unit = unit
+
+    def value(self, ctx: Context) -> Any:
+        """Symbolic value, shape (dim, 1)."""
+        raise NotImplementedError
+
+    def params(self) -> dict[str, Param]:
+        """Params owned by this coordinate, by local name."""
+        return {}
+
+    def children(self) -> tuple[Coordinate, ...]:
+        """Coordinates this one is built from."""
+        return ()
+
+    def __sub__(self, other: Any) -> Coordinate:
+        from .ops import Difference
+
+        return Difference(self, as_coordinate(other, self))
+
+    def __rsub__(self, other: Any) -> Coordinate:
+        from .ops import Difference
+
+        return Difference(as_coordinate(other, self), self)
+
+    def __getitem__(self, index: Any) -> Coordinate:
+        from .ops import Slice
+
+        return Slice(self, index)
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(dim={self.dim}, unit={self.unit!r})"
+
+
+def as_coordinate(x: Any, like: Coordinate) -> Coordinate:
+    """``x`` itself if it is a coordinate, else a live reference shaped like ``like``."""
+    if isinstance(x, Coordinate):
+        return x
+    from .references import Ref
+
+    return Ref("ref", value=x, unit=like.unit)
+
+
+def walk(coord: Coordinate) -> Iterator[Coordinate]:
+    """Every coordinate in the tree of ``coord``, each once, parents first."""
+    seen: set[int] = set()
+    stack = [coord]
+    while stack:
+        node = stack.pop()
+        if id(node) in seen:
+            continue
+        seen.add(id(node))
+        yield node
+        stack.extend(reversed(node.children()))

virtualmodelcontrol/mechanisms/coordinates/frames.py

@@ -0,0 +1,62 @@
+"""Points fixed to a model's frames, at a named site or at an arc parameter s."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import casadi as ca
+
+from ...core.params import Param, as_param
+from .base import Context, Coordinate
+
+
+class FramePoint(Coordinate):
+    """Position [m] of a point fixed to a model frame: a named site, or arc parameter ``s``.
+
+    ``s`` and ``offset`` (in the frame, [m]) become ``episode`` Params, folded in at compile time
+    unless made live. Give both ``at`` and ``s`` for a continuous part of an assembly.
+    """
+
+    def __init__(self, model: Any, at: str | None = None, *, s: Any = None, offset: Any = None):
+        if at is None and s is None:
+            raise ValueError("give a site name `at`, an arc parameter `s`, or both")
+        super().__init__(3, "m")
+        self.model = model
+        self.site = at
+        self.s = None if s is None else as_param(s, "s", bounds=(0.0, 1.0), scope="episode")
+        self.offset = (
+            None if offset is None else as_param(offset, "offset", unit="m", scope="episode")
+        )
+
+    def params(self) -> dict[str, Param]:
+        """The arc parameter and the offset, when given."""
+        out = {}
+        if self.s is not None:
+            out["s"] = self.s
+        if self.offset is not None:
+            out["offset"] = self.offset
+        return out
+
+    def value(self, ctx: Context) -> Any:
+        """Point position in the model's base frame."""
+        if self.s is None:
+            at: Any = self.site
+        elif self.site is None:
+            at = ctx.param(self.s)
+        else:
+            at = (self.site, ctx.param(self.s))
+        R, p = self.model.frame(ctx.q, at, ctx.view(self.model.params))
+        if self.offset is not None:
+            p = p + ca.mtimes(R, ca.reshape(ctx.param(self.offset), 3, 1))
+        return p
+
+    def __repr__(self) -> str:
+        where = ", ".join(
+            x
+            for x in (
+                repr(self.site) if self.site else "",
+                f"s={self.s.value.tolist()}" if self.s is not None else "",
+            )
+            if x
+        )
+        return f"FramePoint({where})"

virtualmodelcontrol/mechanisms/coordinates/joints.py

@@ -0,0 +1,51 @@
+"""Joint coordinates of the robot, and states of virtual degrees of freedom."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import casadi as ca
+import numpy as np
+from numpy.typing import ArrayLike
+
+from .base import Context, Coordinate
+
+
+class Joint(Coordinate):
+    """Entries of the robot's generalized coordinates q: an index, a slice or a list of indices."""
+
+    def __init__(self, index: int | slice | Sequence[int], unit: str = "") -> None:
+        if isinstance(index, slice):
+            if index.stop is None:
+                raise ValueError("a Joint slice needs an explicit stop")
+            indices = list(range(index.stop))[index]
+        elif isinstance(index, int):
+            indices = [index]
+        else:
+            indices = [int(i) for i in index]
+        super().__init__(len(indices), unit)
+        self.indices = indices
+
+    def value(self, ctx: Context) -> Any:
+        """Selected entries of q."""
+        return ca.vertcat(*[ctx.q[i] for i in self.indices])
+
+    def __repr__(self) -> str:
+        return f"Joint({self.indices}, unit={self.unit!r})"
+
+
+class State(Coordinate):
+    """Position of a virtual degree of freedom owned by a controller (part of its state z)."""
+
+    def __init__(self, name: str, dim: int = 1, unit: str = "", initial: ArrayLike = 0.0) -> None:
+        super().__init__(dim, unit)
+        self.name = name
+        self.initial = np.broadcast_to(np.asarray(initial, dtype=float), (dim,)).copy()
+
+    def value(self, ctx: Context) -> Any:
+        """Slice of z."""
+        return ctx.state(self)
+
+    def __repr__(self) -> str:
+        return f"State({self.name!r}, dim={self.dim}, unit={self.unit!r})"