virtualmodelcontrol 0.1.0__py3-none-any.whl

virtualmodelcontrol/__init__.py

@@ -0,0 +1,75 @@
+"""Virtual Model Control: robots and controllers as mechanisms of springs, dampers, inertances."""
+
+from . import sim
+from .compiler import Compiled, compile
+from .control import VMCController
+from .core import SO2, Euclidean, Param, ParamSet, Product, Signals, register
+from .dynamics import Dynamics, compile_dynamics
+from .mechanisms import (
+    Custom,
+    ForceSource,
+    FramePoint,
+    GaussianSpring,
+    Gravity,
+    GravityCompensation,
+    Inertance,
+    Joint,
+    LimitSpring,
+    LinearDamper,
+    LinearSpring,
+    Mechanism,
+    Norm,
+    PointMass,
+    PolynomialSpring,
+    Projection,
+    Ref,
+    SigmoidSpring,
+    Stack,
+    TanhDamper,
+    TanhSpring,
+)
+from .system import VirtualMechanismSystem
+
+try:
+    from ._version import __version__
+except ImportError:  # a source tree that was never installed
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "SO2",
+    "Compiled",
+    "Custom",
+    "Dynamics",
+    "Euclidean",
+    "ForceSource",
+    "FramePoint",
+    "GaussianSpring",
+    "Gravity",
+    "GravityCompensation",
+    "Inertance",
+    "Joint",
+    "LimitSpring",
+    "LinearDamper",
+    "LinearSpring",
+    "Mechanism",
+    "Norm",
+    "Param",
+    "ParamSet",
+    "PointMass",
+    "PolynomialSpring",
+    "Product",
+    "Projection",
+    "Ref",
+    "SigmoidSpring",
+    "Signals",
+    "Stack",
+    "TanhDamper",
+    "TanhSpring",
+    "VMCController",
+    "VirtualMechanismSystem",
+    "__version__",
+    "compile",
+    "compile_dynamics",
+    "register",
+    "sim",
+]

virtualmodelcontrol/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

virtualmodelcontrol/compiler.py

@@ -0,0 +1,158 @@
+"""Compile a virtual mechanism system into CasADi functions: the control law and its energetics."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from .core.params import Binding, ParamSet
+from .mechanisms.coordinates.base import Context
+from .system import VirtualMechanismSystem
+
+ARGS = ["q", "v", "z", "p", "t"]
+OPTS = {"cse": True}  # merges the repeated derivatives of shared coordinates
+
+
+@dataclass
+class Compiled:
+    """CasADi functions of a compiled system.
+
+    ``law`` (→ u, ż), ``tau``, ``energy`` (→ V, T), ``power`` (→ port, dissipation, source) and
+    ``forces`` take (q, v, z, p, t): configuration, velocity, virtual state z = [positions,
+    velocities], live Params p and time t [s]. ``fast`` maps one vector [θ, θ̇, z, p, t] of motor
+    angles and rates to [u, ż]; ``fast_energy`` maps it to V + T.
+    """
+
+    system: VirtualMechanismSystem
+    params: ParamSet
+    live: list[str]
+    law: ca.Function
+    tau: ca.Function
+    energy: ca.Function
+    power: ca.Function
+    forces: ca.Function
+    fast: ca.Function
+    fast_energy: ca.Function
+    component_names: list[str]
+    z0: np.ndarray
+    n_motors: tuple[int, int]
+    n_u: int
+
+    def live_values(self) -> np.ndarray:
+        """Current values of the live Params, packed like p."""
+        return self.params.vector(self.live)
+
+    def live_slices(self) -> dict[str, slice]:
+        """Where each live Param sits in p."""
+        out, offset = {}, 0
+        for name in self.live:
+            size = self.params[name].size
+            out[name] = slice(offset, offset + size)
+            offset += size
+        return out
+
+
+def compile(system: VirtualMechanismSystem, runtime: Iterable[str] = ()) -> Compiled:
+    """Compile ``system`` into CasADi functions.
+
+    ``stage`` Params, and those matching the glob patterns in ``runtime``, stay live inputs; every
+    other Param is folded in at its current value (compile again after changing one).
+    """
+    params = system.params
+    live = params.select(patterns=runtime, scopes=["stage"])
+    binding = Binding(params, live)
+    space, actuation = system.robot.model.space, system.actuation
+    states = system.states
+    nz = sum(s.dim for s in states)
+    q, v, t = ca.SX.sym("q", space.nq), ca.SX.sym("v", space.nv), ca.SX.sym("t")
+    z = ca.SX.sym("z", 2 * nz)
+    zpos, zvel = z[:nz], z[nz:]
+    slices, offset = {}, 0
+    for state in states:
+        slices[id(state)] = slice(offset, offset + state.dim)
+        offset += state.dim
+    ctx = Context(q, binding, z=zpos, t=t, states=slices)
+    G = space.velocity_map(q)
+
+    tau = ca.SX.zeros(space.nv, 1)
+    fz = ca.SX.zeros(nz, 1)
+    Mz = ca.SX.zeros(nz, nz)
+    V, P_diss, P_src = ca.SX(0), ca.SX(0), ca.SX(0)
+    per_component: list[Any] = []
+    names: list[str] = []
+    for name, comp in system.components:
+        y = ctx.value(comp.coord)
+        Jq = ca.mtimes(ca.jacobian(y, q), G)
+        Jz = ca.jacobian(y, zpos)
+        if comp.kind == "inertance":
+            if ca.depends_on(y, ca.vertcat(q, t)):
+                raise ValueError(f"{name}: a controller's inertances may act on its states only")
+            Mz += ca.mtimes([Jz.T, comp.inertance(ctx, y), Jz])
+            continue
+        yd = ca.mtimes(Jq, v) + ca.mtimes(Jz, zvel) + ca.jacobian(y, t)
+        f = comp.force(ctx, y, yd)
+        tau_k = ca.mtimes(Jq.T, f)
+        tau += tau_k
+        fz += ca.mtimes(Jz.T, f)
+        if comp.kind == "storage":
+            V += comp.energy(ctx, y)
+        elif comp.kind == "dissipation":
+            P_diss += ca.dot(f, yd)
+        else:
+            P_src += ca.dot(f, yd)
+        per_component += [y, yd, f, tau_k]
+        names.append(name)
+
+    z0 = np.concatenate([s.initial for s in states] + [np.zeros(nz)])
+    T = 0.5 * ca.dot(zvel, ca.mtimes(Mz, zvel))
+    if nz:
+        M0 = np.array(ca.Function("M", [z, binding.p], [Mz])(z0, binding.values()))
+        if np.linalg.matrix_rank(M0) < nz:
+            raise ValueError("every virtual state needs an inertance (the state mass is singular)")
+        coriolis = ca.jtimes(ca.mtimes(Mz, zvel), zpos, zvel) - ca.gradient(T, zpos)
+        zdot = ca.vertcat(zvel, ca.solve(Mz, fz - coriolis))
+    else:
+        zdot = ca.SX.zeros(0, 1)
+    u = actuation.allocate(tau, q, binding.view(actuation.params))
+
+    args = [q, v, z, binding.p, t]
+    law = ca.Function("law", args, [u, zdot], ARGS, ["u", "zdot"], OPTS)
+    energy = ca.Function("energy", args, [V, T], ARGS, ["V", "T"], OPTS)
+
+    # Fast path: motor angles and rates in, through the exact inverse of the transmission.
+    n_angles, n_rates = actuation.motor_sizes(space)
+    theta, theta_dot = ca.SX.sym("theta", n_angles), ca.SX.sym("theta_dot", n_rates)
+    pa = binding.view(actuation.params)
+    qa = actuation.config_from_motors(theta, pa)
+    va = actuation.velocity_from_motors(qa, theta_dot, pa)
+    x = ca.vertcat(theta, theta_dot, z, binding.p, t)
+    ua, zdot_a = law(qa, va, z, binding.p, t)
+    Va, Ta = energy(qa, va, z, binding.p, t)
+
+    return Compiled(
+        system=system,
+        params=params,
+        live=binding.live,
+        law=law,
+        tau=ca.Function("tau", args, [tau], ARGS, ["tau"], OPTS),
+        energy=energy,
+        power=ca.Function(
+            "power",
+            args,
+            [ca.dot(tau, v), P_diss, P_src],
+            ARGS,
+            ["port", "dissipation", "source"],
+            OPTS,
+        ),
+        forces=ca.Function("forces", args, per_component, OPTS),
+        fast=ca.Function("fast", [x], [ca.vertcat(ua, zdot_a)], ["x"], ["out"], OPTS),
+        fast_energy=ca.Function("fast_energy", [x], [Va + Ta], ["x"], ["E"], OPTS),
+        component_names=names,
+        z0=z0,
+        n_motors=(n_angles, n_rates),
+        n_u=int(u.numel()),
+    )

virtualmodelcontrol/control/__init__.py

@@ -0,0 +1,5 @@
+"""Control: controllers built from compiled virtual mechanisms."""
+
+from .controller import VMCController
+
+__all__ = ["VMCController"]

virtualmodelcontrol/control/controller.py

@@ -0,0 +1,76 @@
+"""VMCController: runs a compiled virtual mechanism step by step."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ..core.signals import Signals
+
+
+class VMCController:
+    """Runs a compiled system: motor angles and rates in, motor torques out.
+
+    ``step`` reads ``motor_position`` [rad] and ``motor_velocity`` [rad/s] and returns
+    ``motor_torque`` [N·m]. Virtual states are integrated by semi-implicit Euler over the measured
+    dt. Live Params start at their values when compiled; ``set`` changes them here only.
+    """
+
+    def __init__(self, compiled: Any) -> None:
+        self.compiled = compiled
+        self.params = compiled.live_values()
+        self._slices = compiled.live_slices()
+        self.z = compiled.z0.copy()
+        self.t: float | None = None
+        self._x: np.ndarray | None = None
+
+    def reset(self, t: float, meas: Signals | None = None, z0: ArrayLike | None = None) -> None:
+        """Restart at time ``t`` [s] from the initial virtual state (or ``z0``)."""
+        self.z = self.compiled.z0.copy() if z0 is None else np.array(z0, dtype=float)
+        self._x = None if meas is None else self._pack(meas, t)
+        self.t = t
+
+    def step(self, t: float, meas: Signals) -> Signals:
+        """One control step at time ``t`` [s]; returns the motor torques."""
+        dt = 0.0 if self.t is None else t - self.t
+        x = self._pack(meas, t)
+        out = np.asarray(self.compiled.fast(x)).ravel()
+        nu, nz = self.compiled.n_u, self.z.size // 2
+        if nz:
+            self.z[nz:] += dt * out[nu + nz :]
+            self.z[:nz] += dt * self.z[nz:]
+        self.t, self._x = t, x
+        return Signals(t, motor_torque=out[:nu])
+
+    def set(self, values: Mapping[str, ArrayLike] | None = None, **kwargs: ArrayLike) -> float:
+        """Change live Params at once; returns the exact jump of the controller's energy [J]."""
+        new = self.params.copy()
+        for name, value in {**(values or {}), **kwargs}.items():
+            if name not in self._slices:
+                raise KeyError(
+                    f"{name!r} is not a live Param here; compile with runtime=[{name!r}] to change "
+                    "it while running"
+                )
+            new[self._slices[name]] = np.ravel(value, order="F")
+        jump = 0.0 if self._x is None else self._energy(new) - self._energy(self.params)
+        self.params = new
+        return jump
+
+    def energy(self) -> float:
+        """Energy of the controller at the last step [J] (stored plus virtual kinetic)."""
+        return 0.0 if self._x is None else self._energy(self.params)
+
+    def _energy(self, params: np.ndarray) -> float:
+        assert self._x is not None
+        x = self._x.copy()
+        n = self._x.size - 1 - params.size
+        x[n:-1] = params
+        return float(self.compiled.fast_energy(x))
+
+    def _pack(self, meas: Signals, t: float) -> np.ndarray:
+        return np.concatenate(
+            [meas["motor_position"], meas["motor_velocity"], self.z, self.params, [t]]
+        )

virtualmodelcontrol/core/__init__.py

@@ -0,0 +1,24 @@
+"""Core: parameters, spaces, symbolic helpers, signals, registry and units."""
+
+from .params import SCOPES, Binding, Param, ParamSet, as_param, constants
+from .registry import get, load_plugins, names, register
+from .signals import Signals
+from .space import SO2, Euclidean, Product, Space
+
+__all__ = [
+    "SCOPES",
+    "SO2",
+    "Binding",
+    "Euclidean",
+    "Param",
+    "ParamSet",
+    "Product",
+    "Signals",
+    "Space",
+    "as_param",
+    "constants",
+    "get",
+    "load_plugins",
+    "names",
+    "register",
+]

virtualmodelcontrol/core/params.py

@@ -0,0 +1,244 @@
+"""Parameters: named numbers with units, bounds and scopes, packed into one flat vector."""
+
+from __future__ import annotations
+
+import fnmatch
+from collections.abc import Iterable, Iterator, Mapping
+from typing import Any, Literal
+
+import casadi as ca
+import numpy as np
+from numpy.typing import ArrayLike
+
+Scope = Literal["fixed", "design", "episode", "stage"]
+SCOPES: tuple[str, ...] = ("fixed", "design", "episode", "stage")
+
+
+class Param:
+    """A named number (scalar, vector or matrix) with a unit, bounds and a scope.
+
+    The scope is the time scale on which the value may change: ``fixed`` (never), ``design``
+    (with the hardware), ``episode`` (between runs) or ``stage`` (at every step).
+    """
+
+    __slots__ = ("_value", "bounds", "name", "scale", "scope", "unit")
+
+    def __init__(
+        self,
+        name: str,
+        value: ArrayLike,
+        *,
+        unit: str = "",
+        bounds: tuple[Any, Any] = (-np.inf, np.inf),
+        scale: float = 1.0,
+        scope: Scope = "fixed",
+    ) -> None:
+        if scope not in SCOPES:
+            raise ValueError(f"scope must be one of {SCOPES}, got {scope!r}")
+        value = np.array(value, dtype=float)
+        if value.ndim > 2:
+            raise ValueError(f"Param {name!r} must be a scalar, vector or matrix")
+        self.name = name
+        self._value = value
+        self.unit = unit
+        self.bounds = (bounds[0], bounds[1])
+        self.scale = float(scale)
+        self.scope = scope
+
+    @property
+    def value(self) -> np.ndarray:
+        """Current value, as a float array of fixed shape."""
+        return self._value
+
+    @value.setter
+    def value(self, value: ArrayLike) -> None:
+        new = np.array(value, dtype=float)
+        if new.shape != self._value.shape:
+            if new.size != self._value.size:
+                raise ValueError(
+                    f"Param {self.name!r} has shape {self._value.shape}, got a value of shape "
+                    f"{new.shape}"
+                )
+            new = new.reshape(self._value.shape)
+        self._value = new
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Shape of the value: () for a scalar, (n,) for a vector, (r, c) for a matrix."""
+        return self._value.shape
+
+    @property
+    def size(self) -> int:
+        """Number of entries in the value."""
+        return self._value.size
+
+    def __repr__(self) -> str:
+        unit = f", unit={self.unit!r}" if self.unit else ""
+        return f"Param({self.name!r}, {self._value.tolist()!r}{unit}, scope={self.scope!r})"
+
+
+def as_param(
+    value: Any,
+    name: str,
+    *,
+    unit: str = "",
+    bounds: tuple[Any, Any] = (-np.inf, np.inf),
+    scope: Scope = "fixed",
+) -> Param:
+    """Return ``value`` if it already is a Param, else wrap it in a new Param."""
+    if isinstance(value, Param):
+        return value
+    return Param(name, value, unit=unit, bounds=bounds, scope=scope)
+
+
+class ParamSet(Mapping[str, Param]):
+    """Ordered, named collection of Params; packs their values into one flat vector.
+
+    A Param object appears once: adding it again under another name keeps the first name.
+    Matrices are packed column by column, as CasADi stores them.
+    """
+
+    def __init__(self, params: Iterable[Param] = ()) -> None:
+        self._params: dict[str, Param] = {}
+        self._names: dict[int, str] = {}
+        for param in params:
+            self.add(param)
+
+    def add(self, param: Param, name: str | None = None, *, rename: bool = False) -> str:
+        """Add ``param`` under ``name`` (default: its own name) and return the name used.
+
+        With ``rename``, a name already taken gets a numeric suffix instead of raising.
+        """
+        if id(param) in self._names:
+            return self._names[id(param)]
+        name = param.name if name is None else name
+        if name in self._params:
+            if not rename:
+                raise ValueError(f"duplicate parameter name {name!r}")
+            k = 2
+            while f"{name}{k}" in self._params:
+                k += 1
+            name = f"{name}{k}"
+        self._params[name] = param
+        self._names[id(param)] = name
+        return name
+
+    def merge(self, other: ParamSet, prefix: str = "") -> None:
+        """Add every Param of ``other``, with names prefixed by ``prefix.``."""
+        for name, param in other.items():
+            self.add(param, f"{prefix}.{name}" if prefix else name)
+
+    def __getitem__(self, name: str) -> Param:
+        try:
+            return self._params[name]
+        except KeyError:
+            raise KeyError(f"no parameter {name!r}; known: {list(self._params)}") from None
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._params)
+
+    def __len__(self) -> int:
+        return len(self._params)
+
+    def __repr__(self) -> str:
+        return f"ParamSet({list(self._params)})"
+
+    def name_of(self, param: Param) -> str:
+        """Name under which ``param`` is stored."""
+        return self._names[id(param)]
+
+    def has(self, param: Param) -> bool:
+        """True if this exact Param object is in the set."""
+        return id(param) in self._names
+
+    def select(self, patterns: Iterable[str] = (), scopes: Iterable[str] = ()) -> list[str]:
+        """Names matching any glob pattern (``ctrl.*.stiffness``) or having any of the scopes."""
+        patterns, scopes = list(patterns), set(scopes)
+        return [
+            name
+            for name, param in self._params.items()
+            if param.scope in scopes or any(fnmatch.fnmatchcase(name, pat) for pat in patterns)
+        ]
+
+    def size(self, names: Iterable[str] | None = None) -> int:
+        """Total number of entries of the named Params (default: all)."""
+        return sum(self[n].size for n in self._names_or_all(names))
+
+    def vector(self, names: Iterable[str] | None = None) -> np.ndarray:
+        """Values of the named Params (default: all) packed into one flat vector."""
+        parts = [self[n].value.ravel(order="F") for n in self._names_or_all(names)]
+        return np.concatenate(parts) if parts else np.zeros(0)
+
+    def set_vector(self, x: ArrayLike, names: Iterable[str] | None = None) -> None:
+        """Unpack a flat vector into the named Params (default: all)."""
+        x = np.asarray(x, dtype=float).ravel()
+        names = self._names_or_all(names)
+        if x.size != self.size(names):
+            raise ValueError(f"expected {self.size(names)} values, got {x.size}")
+        offset = 0
+        for n in names:
+            param = self[n]
+            param.value = x[offset : offset + param.size].reshape(param.shape, order="F")
+            offset += param.size
+
+    def to_dict(self) -> dict[str, Any]:
+        """Values by name, as plain floats and lists."""
+        return {name: param.value.tolist() for name, param in self._params.items()}
+
+    def update(self, values: Mapping[str, ArrayLike]) -> None:
+        """Set values by name."""
+        for name, value in values.items():
+            self[name].value = value
+
+    def _names_or_all(self, names: Iterable[str] | None) -> list[str]:
+        return list(self._params) if names is None else list(names)
+
+
+def _shaped(block: Any, shape: tuple[int, ...]) -> Any:
+    return ca.reshape(block, shape[0], shape[1]) if len(shape) == 2 else block
+
+
+class Binding:
+    """CasADi expressions for the Params of a set.
+
+    Live Params are slices of one symbol ``p``; the others are folded in as constants at their
+    current values.
+    """
+
+    def __init__(self, params: ParamSet, live: Iterable[str] = (), symbol: Any = ca.SX) -> None:
+        live = set(live)
+        unknown = live - set(params)
+        if unknown:
+            raise KeyError(f"unknown live parameters {sorted(unknown)}")
+        self.params = params
+        self.live = [name for name in params if name in live]
+        self.p = symbol.sym("p", params.size(self.live))
+        self._expr: dict[int, Any] = {}
+        offset = 0
+        for name, param in params.items():
+            if name in live:
+                block = self.p[offset : offset + param.size]
+                offset += param.size
+            else:
+                block = ca.DM(param.value.ravel(order="F"))
+            self._expr[id(param)] = _shaped(block, param.shape)
+
+    def __call__(self, param: Param) -> Any:
+        """Expression of ``param``."""
+        try:
+            return self._expr[id(param)]
+        except KeyError:
+            raise KeyError(f"{param!r} is not in this binding") from None
+
+    def view(self, params: Mapping[str, Param]) -> dict[str, Any]:
+        """Expressions of a model's Params, keyed by the model's own names."""
+        return {name: self(param) for name, param in params.items()}
+
+    def values(self) -> np.ndarray:
+        """Current values of the live Params, packed like ``p``."""
+        return self.params.vector(self.live)
+
+
+def constants(params: Mapping[str, Param]) -> dict[str, Any]:
+    """Current values of Params as CasADi constants, keyed by name (a numeric view)."""
+    return {name: _shaped(ca.DM(p.value.ravel(order="F")), p.shape) for name, p in params.items()}

virtualmodelcontrol/core/registry.py

@@ -0,0 +1,54 @@
+"""Registry: names usable from configuration files, extended by plugins through entry points."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from importlib.metadata import entry_points
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+GROUP = "virtualmodelcontrol.plugins"
+"""Entry-point group scanned for plugins; loading an entry point runs its ``register`` calls."""
+
+_REGISTRY: dict[str, dict[str, Any]] = {}
+_plugins_loaded = False
+
+
+def register(kind: str, name: str) -> Callable[[T], T]:
+    """Class or function decorator: make ``obj`` available as ``get(kind, name)``."""
+
+    def decorate(obj: T) -> T:
+        table = _REGISTRY.setdefault(kind, {})
+        if name in table and table[name] is not obj:
+            raise ValueError(f"a {kind} named {name!r} is already registered")
+        table[name] = obj
+        return obj
+
+    return decorate
+
+
+def get(kind: str, name: str) -> Any:
+    """Registered object; plugins are loaded on the first miss."""
+    if name not in _REGISTRY.get(kind, {}):
+        load_plugins()
+    table = _REGISTRY.get(kind, {})
+    if name not in table:
+        raise KeyError(f"no {kind} named {name!r}; known: {sorted(table)}")
+    return table[name]
+
+
+def names(kind: str) -> list[str]:
+    """Registered names of one kind, plugins included."""
+    load_plugins()
+    return sorted(_REGISTRY.get(kind, {}))
+
+
+def load_plugins() -> None:
+    """Import every entry point in the group ``virtualmodelcontrol.plugins`` (once)."""
+    global _plugins_loaded
+    if _plugins_loaded:
+        return
+    _plugins_loaded = True
+    for ep in entry_points(group=GROUP):
+        ep.load()