cozyclay 2.1.1 → 3.0.0

package/tools/track/py/rig.py

@@ -1,508 +0,0 @@
-"""Differentiable Studio rig for the known-character tracker (plan mocap-rearch todo 6).
-
-Two deliberately separate layers:
-
-1. `cskel27_fk(state, bone_scale, rig)` -> posedJoints. cskel27 forward kinematics
-   over the canonical neutral skeleton grown by `bone_scale` - exactly the
-   quantity a motion NPZ stores as `posed_joints` (tools/ardy/smpl-cskel27.mjs,
-   tools/bench/fit/motion.mjs regenerateJoints) and the scorers read. Keypoint and
-   contact terms use these joints.
-2. `studio_skin(globals, posed_joints, rig, ...)` -> bone world transforms ->
-   vertices. A torch port of Studio playback's positional skinning
-   (src/ardy/playback.js applyMotionFrame, prepOf, scaledOffsets, boneStretch)
-   followed by linear-blend skinning of the exported rig mesh. Silhouette and box
-   terms use these vertices.
-
-State convention (NPZ `local_rot_mats` convention): per frame
-    transl  (T, 3)      root/pelvis world position = posed_joints[:, 0]
-    root6d  (T, 6)      a_t, the root rotation = local_rot_mats[:, 0]
-    local6d (T, 26, 6)  u_t,j for cskel27 joints 1..26 (local to the parent)
-Zero pose = identity locals = the neutral T-pose, so 6D vectors are the rotations
-themselves (no rest delta). `rot6d_to_matrix` follows the ultrabrain design
-section 1: the 6D vector is the first two COLUMNS of the matrix, Gram-Schmidt.
-
-Rig input: the npz written by tools/track/export-rig.mjs (layout and units in
-tools/track/rig-dump.mjs). `load_rig` validates it (`RigError.code`:
-`bad-rig`, `bad-skin-weights`) and decimates the merged mesh to <= 8,000 faces,
-carrying skin weights from the nearest original vertex (renormalised).
-
-World frame: Studio world, metres, +Y up. `studio_skin` returns the rig as Studio
-places it under a Character group at the origin with no clip-to-scene yaw:
-    rig-space bone = s * (posed - anchor_xz) + R_global @ offset
-with s = prep scale (rig leg height / canonical leg height, rig units) and the
-0.01 rig root scale taking rig units to metres. Pass `anchor_xz=None` for the
-tracker (no anchoring: vertices live in the posedJoints frame, scaled by the
-rig's leg ratio about the XZ origin exactly as Studio scales them about the
-anchor), or the anchor frame's root XZ to reproduce Studio playback bit-for-bit.
-Nothing here optimises bone lengths, shape or camera.
-"""
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass
-from pathlib import Path
-
-import numpy as np
-import torch
-
-J = 27
-INFLUENCES = 4
-WEIGHT_TOLERANCE = 1e-5  # same as tools/track/rig-dump.mjs
-MAX_FACES = 8000
-EPS = 1e-8
-
-REQUIRED_MEMBERS = (
-    "meta_json", "rig_root_scale", "vertices", "faces", "skin_index", "skin_weight",
-    "bone_parent", "bone_rest_local", "bone_bind_inverse",
-    "prep_bone", "prep_scale", "prep_offsets", "prep_bind_pos", "prep_bind_quat",
-    "prep_bind_scale", "prep_bind_local_pos", "prep_parent_bind_world",
-    "prep_chain_parent", "prep_chain_rel", "prep_canonical_bone_length",
-    "prep_rig_bone_length", "prep_stretched_leaves", "prep_stretched_leaf_local_pos",
-    "prep_stretched_leaf_local_quat", "hierarchy_preserved", "girdle_joints",
-    "ardy_neutral_min_y", "cskel27_parents", "cskel27_neutral",
-)
-
-
-class RigError(ValueError):
-    """Malformed rig input. `code` is the machine-readable reason."""
-
-    def __init__(self, code: str, detail: str):
-        super().__init__(f"{code}: {detail}")
-        self.code = code
-
-
-# --- rotations -----------------------------------------------------------------
-
-def rot6d_to_matrix(a: torch.Tensor) -> torch.Tensor:
-    """(..., 6) -> (..., 3, 3). a = [a1, a2]; columns b1, b2, b3 (Gram-Schmidt)."""
-    # Clamp (not add) the epsilon: columns stay exactly unit length unless degenerate.
-    a1, a2 = a[..., 0:3], a[..., 3:6]
-    b1 = a1 / a1.norm(dim=-1, keepdim=True).clamp_min(EPS)
-    a2 = a2 - b1 * (b1 * a2).sum(dim=-1, keepdim=True)
-    b2 = a2 / a2.norm(dim=-1, keepdim=True).clamp_min(EPS)
-    b3 = torch.cross(b1, b2, dim=-1)
-    return torch.stack((b1, b2, b3), dim=-1)
-
-
-def matrix_to_rot6d(R: torch.Tensor) -> torch.Tensor:
-    """(..., 3, 3) -> (..., 6): the first two columns."""
-    return torch.cat((R[..., :, 0], R[..., :, 1]), dim=-1)
-
-
-def quat_to_matrix(q: torch.Tensor) -> torch.Tensor:
-    """(..., 4) three.js [x, y, z, w] (normalised here) -> (..., 3, 3)."""
-    q = q / q.norm(dim=-1, keepdim=True)
-    x, y, z, w = q.unbind(-1)
-    return torch.stack((
-        1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y),
-        2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x),
-        2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y),
-    ), dim=-1).reshape(q.shape[:-1] + (3, 3))
-
-
-def _affine(R: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """(..., 3, 3), (..., 3) -> (..., 4, 4)."""
-    top = torch.cat((R, t.unsqueeze(-1)), dim=-1)
-    bottom = torch.zeros(top.shape[:-2] + (1, 4), dtype=top.dtype, device=top.device)
-    bottom[..., 0, 3] = 1
-    return torch.cat((top, bottom), dim=-2)
-
-
-def _apply(M: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
-    """Affine (..., 4, 4) applied to points (..., 3)."""
-    return (M[..., :3, :3] @ p.unsqueeze(-1)).squeeze(-1) + M[..., :3, 3]
-
-
-# --- rig -----------------------------------------------------------------------
-
-@dataclass
-class Rig:
-    """Exported Studio rig as torch tensors (float, on `device`)."""
-
-    meta: dict
-    parents: list  # cskel27 parent per joint, -1 for the root
-    neutral: torch.Tensor  # (27, 3) canonical neutral, hips at the origin
-    min_y: float  # ARDY_NEUTRAL_MIN_Y
-    root_scale: float  # rig units -> metres (0.01)
-    bone_names: list
-    bone_parent: list  # per exported bone, -1 for bones under the (static) rig root
-    bone_rest_local: torch.Tensor  # (B, 4, 4) rig units
-    bone_rest_scale: torch.Tensor  # (B, 3)
-    bind_inverse: torch.Tensor  # (B, 4, 4) metres
-    # prepOf
-    prep_bone: list  # cskel27 -> bone index or -1
-    prep_scale: float
-    prep_offsets: torch.Tensor  # (27, 3)
-    prep_bind_pos: torch.Tensor  # (27, 3)
-    prep_bind_rot: torch.Tensor  # (27, 3, 3)
-    prep_bind_scale: torch.Tensor  # (27, 3)
-    prep_bind_local_pos: torch.Tensor  # (27, 3)
-    prep_parent_bind_world: torch.Tensor  # (27, 4, 4)
-    prep_chain_parent: list
-    prep_chain_rel: torch.Tensor  # (27, 4, 4)
-    canonical_bone_length: np.ndarray  # (27,)
-    rig_bone_length: np.ndarray  # (27,)
-    leaves: list  # [(bone, joint, local_pos (3,), local_rot (3, 3))]
-    hierarchy_preserved: np.ndarray  # (27,) bool
-    girdle: np.ndarray  # (27,) bool
-    # full-resolution merged mesh (rest pose, metres)
-    vertices: torch.Tensor  # (V, 3)
-    skin_index: torch.Tensor  # (V, 4) long
-    skin_weight: torch.Tensor  # (V, 4)
-    # decimated mesh used by the tracker's silhouette/box terms
-    dec_vertices: torch.Tensor  # (Vd, 3)
-    dec_faces: torch.Tensor  # (Fd, 3) long
-    dec_skin_index: torch.Tensor  # (Vd, 4) long
-    dec_skin_weight: torch.Tensor  # (Vd, 4)
-    dec_weights: torch.Tensor  # (Vd, B) dense blend weights
-    dec_source: np.ndarray  # (Vd,) nearest original vertex
-
-
-def _require(z, name, shape=None, dtype_kind=None):
-    if name not in z:
-        raise RigError("bad-rig", f"missing member {name}")
-    value = np.asarray(z[name])
-    if shape is not None and (value.ndim != len(shape) or any(s is not None and s != d for s, d in zip(shape, value.shape))):
-        raise RigError("bad-rig", f"{name} has shape {value.shape}, expected {shape}")
-    if dtype_kind == "f" and not np.all(np.isfinite(value)):
-        raise RigError("bad-rig", f"{name} has non-finite values")
-    return value
-
-
-def validate_skin_weights(skin_weight: np.ndarray, tol: float = WEIGHT_TOLERANCE) -> None:
-    """Raise `bad-skin-weights` unless every row is non-negative and sums to 1."""
-    w = np.asarray(skin_weight, dtype=np.float64)
-    if not np.all(np.isfinite(w)):
-        raise RigError("bad-skin-weights", "non-finite skin weight")
-    if (w < 0).any():
-        v = int(np.argwhere(w < 0)[0, 0])
-        raise RigError("bad-skin-weights", f"vertex {v} has a negative weight")
-    err = np.abs(w.sum(axis=1) - 1)
-    if (err > tol).any():
-        v = int(err.argmax())
-        raise RigError("bad-skin-weights", f"vertex {v} weights sum to {w[v].sum():.7f}")
-
-
-def decimate(vertices: np.ndarray, faces: np.ndarray, skin_index: np.ndarray, skin_weight: np.ndarray,
-             max_faces: int = MAX_FACES):
-    """Weld, simplify to <= max_faces, transfer skin weights from the nearest original vertex.
-
-    Returns (vertices (Vd,3), faces (Fd,3), skin_index (Vd,4), skin_weight (Vd,4), source (Vd,)).
-    The export is unindexed per triangle, so identical positions are welded first
-    (fast_simplification needs shared vertices to collapse edges).
-    """
-    import fast_simplification
-    from scipy.spatial import cKDTree
-
-    welded, first, inverse = np.unique(vertices, axis=0, return_index=True, return_inverse=True)
-    inverse = inverse.reshape(-1)
-    wf = inverse[faces]
-    keep = (wf[:, 0] != wf[:, 1]) & (wf[:, 1] != wf[:, 2]) & (wf[:, 0] != wf[:, 2])
-    wf = wf[keep]
-    reduction = 1.0 - max_faces / len(wf) if len(wf) > max_faces else 0.0
-    for _ in range(20):
-        if reduction <= 0:
-            points, out_faces = welded, wf
-        else:
-            points, out_faces = fast_simplification.simplify(
-                welded.astype(np.float32), wf.astype(np.int32), target_reduction=min(reduction, 0.999))
-        if len(out_faces) <= max_faces:
-            break
-        reduction += (1 - reduction) * 0.05
-    else:
-        raise RigError("bad-rig", f"could not decimate below {max_faces} faces")
-    points = np.asarray(points, dtype=np.float32)
-    out_faces = np.asarray(out_faces, dtype=np.int64)
-    used = np.unique(out_faces)  # drop vertices no face references
-    remap = np.full(len(points), -1, dtype=np.int64)
-    remap[used] = np.arange(len(used))
-    points, out_faces = points[used], remap[out_faces]
-    _, nearest = cKDTree(vertices).query(points, k=1)
-    source = np.asarray(nearest, dtype=np.int64)
-    w = skin_weight[source].astype(np.float64)
-    w = w / w.sum(axis=1, keepdims=True)
-    return points, out_faces, skin_index[source].astype(np.int64), w.astype(np.float32), source
-
-
-def load_rig(path, device="cpu", dtype=torch.float32, max_faces: int = MAX_FACES) -> Rig:
-    path = Path(path)
-    if not path.is_file():
-        raise RigError("bad-rig", f"rig npz not found: {path}")
-    with np.load(path) as z:
-        missing = [name for name in REQUIRED_MEMBERS if name not in z]
-        if missing:
-            raise RigError("bad-rig", f"missing members: {', '.join(missing)}")
-        meta = json.loads(bytes(z["meta_json"].astype(np.uint8)))
-        vertices = _require(z, "vertices", (None, 3), "f").astype(np.float32)
-        V = len(vertices)
-        faces = _require(z, "faces", (None, 3)).astype(np.int64)
-        skin_index = _require(z, "skin_index", (V, INFLUENCES)).astype(np.int64)
-        skin_weight = _require(z, "skin_weight", (V, INFLUENCES), "f").astype(np.float32)
-        bone_parent = _require(z, "bone_parent", (None,)).astype(np.int64)
-        B = len(bone_parent)
-        arrays = {name: _require(z, name, shape, "f") for name, shape in (
-            ("bone_rest_local", (B, 4, 4)), ("bone_bind_inverse", (B, 4, 4)),
-            ("prep_offsets", (J, 3)), ("prep_bind_pos", (J, 3)), ("prep_bind_quat", (J, 4)),
-            ("prep_bind_scale", (J, 3)), ("prep_bind_local_pos", (J, 3)),
-            ("prep_parent_bind_world", (J, 4, 4)), ("prep_chain_rel", (J, 4, 4)),
-            ("prep_canonical_bone_length", (J,)), ("prep_rig_bone_length", (J,)),
-            ("cskel27_neutral", (J, 3)), ("rig_root_scale", (1,)), ("prep_scale", (1,)),
-            ("ardy_neutral_min_y", (1,)),
-        )}
-        prep_bone = _require(z, "prep_bone", (J,)).astype(np.int64)
-        chain_parent = _require(z, "prep_chain_parent", (J,)).astype(np.int64)
-        parents = _require(z, "cskel27_parents", (J,)).astype(np.int64)
-        leaves_idx = _require(z, "prep_stretched_leaves", (None, 2)).astype(np.int64)
-        leaf_pos = _require(z, "prep_stretched_leaf_local_pos", (len(leaves_idx), 3), "f")
-        leaf_quat = _require(z, "prep_stretched_leaf_local_quat", (len(leaves_idx), 4), "f")
-        preserved = _require(z, "hierarchy_preserved", (J,)).astype(bool)
-        girdle = _require(z, "girdle_joints", (J,)).astype(bool)
-
-    if len(meta.get("boneNames", [])) != B:
-        raise RigError("bad-rig", f"meta boneNames has {len(meta.get('boneNames', []))} names for {B} bones")
-    if faces.size and (faces.min() < 0 or faces.max() >= V):
-        raise RigError("bad-rig", "face index out of range")
-    if skin_index.min() < 0 or skin_index.max() >= B:
-        raise RigError("bad-rig", "skin index out of range")
-    for i, p in enumerate(bone_parent):
-        if not -1 <= p < i:
-            raise RigError("bad-rig", f"bone {i} parent {p} does not precede it")
-    for j, p in enumerate(parents):
-        if not (p == -1 if j == 0 else 0 <= p < j):
-            raise RigError("bad-rig", f"cskel27 parent {p} of joint {j} is not topological")
-    for j in range(J):
-        if not -1 <= prep_bone[j] < B:
-            raise RigError("bad-rig", f"prep_bone[{j}] out of range")
-        cp = chain_parent[j]
-        if cp >= 0 and (cp >= j or prep_bone[cp] < 0 or prep_bone[j] < 0):
-            raise RigError("bad-rig", f"prep_chain_parent[{j}]={cp} must be an earlier mapped joint")
-    validate_skin_weights(skin_weight)
-
-    dv, df, dsi, dsw, source = decimate(vertices, faces, skin_index, skin_weight, max_faces)
-    dense = np.zeros((len(dv), B), dtype=np.float32)
-    np.add.at(dense, (np.repeat(np.arange(len(dv)), INFLUENCES), dsi.reshape(-1)), dsw.reshape(-1))
-
-    t = lambda a: torch.as_tensor(np.asarray(a, dtype=np.float32), device=device).to(dtype)  # noqa: E731
-    rest_local = arrays["bone_rest_local"]
-    return Rig(
-        meta=meta,
-        parents=[int(p) for p in parents],
-        neutral=t(arrays["cskel27_neutral"]),
-        min_y=float(arrays["ardy_neutral_min_y"][0]),
-        root_scale=float(arrays["rig_root_scale"][0]),
-        bone_names=list(meta["boneNames"]),
-        bone_parent=[int(p) for p in bone_parent],
-        bone_rest_local=t(rest_local),
-        bone_rest_scale=t(np.linalg.norm(rest_local[:, :3, :3], axis=1)),
-        bind_inverse=t(arrays["bone_bind_inverse"]),
-        prep_bone=[int(b) for b in prep_bone],
-        prep_scale=float(arrays["prep_scale"][0]),
-        prep_offsets=t(arrays["prep_offsets"]),
-        prep_bind_pos=t(arrays["prep_bind_pos"]),
-        prep_bind_rot=quat_to_matrix(t(arrays["prep_bind_quat"])),
-        prep_bind_scale=t(arrays["prep_bind_scale"]),
-        prep_bind_local_pos=t(arrays["prep_bind_local_pos"]),
-        prep_parent_bind_world=t(arrays["prep_parent_bind_world"]),
-        prep_chain_parent=[int(c) for c in chain_parent],
-        prep_chain_rel=t(arrays["prep_chain_rel"]),
-        canonical_bone_length=np.asarray(arrays["prep_canonical_bone_length"], dtype=np.float64),
-        rig_bone_length=np.asarray(arrays["prep_rig_bone_length"], dtype=np.float64),
-        leaves=[(int(b), int(j), t(p), quat_to_matrix(t(q))) for (b, j), p, q in zip(leaves_idx, leaf_pos, leaf_quat)],
-        hierarchy_preserved=preserved,
-        girdle=girdle,
-        vertices=t(vertices),
-        skin_index=torch.as_tensor(skin_index, device=device),
-        skin_weight=t(skin_weight),
-        dec_vertices=t(dv),
-        dec_faces=torch.as_tensor(df, device=device),
-        dec_skin_index=torch.as_tensor(dsi, device=device),
-        dec_skin_weight=t(dsw),
-        dec_weights=t(dense),
-        dec_source=source,
-    )
-
-
-# --- layer 1: cskel27 FK ---------------------------------------------------------
-
-@dataclass
-class State:
-    """Optimisable per-frame pose (see module docstring for the convention)."""
-
-    transl: torch.Tensor  # (T, 3)
-    root6d: torch.Tensor  # (T, 6)
-    local6d: torch.Tensor  # (T, 26, 6), joints 1..26
-
-    @staticmethod
-    def from_motion(local_rot_mats, root_positions, device="cpu", dtype=torch.float32) -> "State":
-        R = torch.as_tensor(np.asarray(local_rot_mats), device=device).to(dtype)
-        x = torch.as_tensor(np.asarray(root_positions), device=device).to(dtype)
-        if R.ndim != 4 or R.shape[1:] != (J, 3, 3) or x.shape != (R.shape[0], 3):
-            raise RigError("bad-motion", f"expected local_rot_mats (T,27,3,3) and root (T,3), got {tuple(R.shape)} {tuple(x.shape)}")
-        return State(x.clone(), matrix_to_rot6d(R[:, 0]), matrix_to_rot6d(R[:, 1:]))
-
-    @staticmethod
-    def identity(frames: int, root=(0.0, 0.0, 0.0), device="cpu", dtype=torch.float32) -> "State":
-        eye6 = torch.tensor([1.0, 0, 0, 0, 1, 0], device=device, dtype=dtype)
-        return State(torch.tensor(root, device=device, dtype=dtype).expand(frames, 3).clone(),
-                     eye6.expand(frames, 6).clone(), eye6.expand(frames, J - 1, 6).clone())
-
-    def local_rot_mats(self) -> torch.Tensor:
-        """(T, 27, 3, 3) in the NPZ local_rot_mats convention."""
-        return torch.cat((rot6d_to_matrix(self.root6d).unsqueeze(1), rot6d_to_matrix(self.local6d)), dim=1)
-
-
-def global_rotations(local: torch.Tensor, parents) -> torch.Tensor:
-    """(T, 27, 3, 3) local -> global rotations down the cskel27 tree."""
-    out = [None] * J
-    for j in range(J):
-        p = parents[j]
-        out[j] = local[:, j] if p < 0 else out[p] @ local[:, j]
-    return torch.stack(out, dim=1)
-
-
-def cskel27_offsets(rig: Rig, bone_scale=None) -> torch.Tensor:
-    """(27, 3) bone offsets in the parent frame: boneScale[j] * (neutral[j] - neutral[parent])."""
-    n = rig.neutral
-    parents = torch.tensor([max(p, 0) for p in rig.parents], device=n.device)
-    offsets = n - n[parents]  # the root row is n[0] - n[0] = 0
-    if bone_scale is not None:
-        offsets = offsets * _bone_scale(bone_scale, n).unsqueeze(-1)
-    return offsets
-
-
-def _bone_scale(bone_scale, like: torch.Tensor) -> torch.Tensor:
-    bs = torch.as_tensor(np.asarray(bone_scale) if not torch.is_tensor(bone_scale) else bone_scale, device=like.device).to(like.dtype)
-    if bs.shape != (J,) or not torch.isfinite(bs).all() or (bs <= 0).any():
-        raise RigError("bad-bone-scale", "bone_scale must be 27 positive finite factors")
-    return bs
-
-
-def cskel27_fk(state: State, bone_scale, rig: Rig):
-    """Layer 1. Returns (posed_joints (T, 27, 3), globals (T, 27, 3, 3)).
-
-    posed_joints is the motion NPZ `posed_joints`: FK over the neutral skeleton
-    grown by `bone_scale` (None == all ones) with the root at `state.transl`.
-    """
-    globals_ = global_rotations(state.local_rot_mats(), rig.parents)
-    offsets = cskel27_offsets(rig, bone_scale)
-    pos = [None] * J
-    for j in range(J):
-        p = rig.parents[j]
-        pos[j] = state.transl if p < 0 else pos[p] + (globals_[:, p] @ offsets[j].unsqueeze(-1)).squeeze(-1)
-    return torch.stack(pos, dim=1), globals_
-
-
-# --- layer 2: Studio positional skinning + LBS ------------------------------------
-
-def _studio_offsets(rig: Rig, bone_scale) -> torch.Tensor:
-    """playback.js prep.offsets, or scaledOffsets(prep, boneScale) for a take with boneScale."""
-    if bone_scale is None:
-        return rig.prep_offsets
-    grown = cskel27_offsets(rig, bone_scale)
-    g = [None] * J
-    for j in range(J):
-        p = rig.parents[j]
-        g[j] = rig.neutral[j] if p < 0 else g[p] + grown[j]
-    g = torch.stack(g)
-    shift = torch.tensor([0.0, rig.min_y, 0.0], dtype=g.dtype, device=g.device)
-    return rig.prep_bind_pos - rig.prep_scale * (g - shift)
-
-
-def _stretch(rig: Rig, bone_scale, j: int):
-    """playback.js boneStretch: 1 without boneScale; girdle joints take the raw factor."""
-    if bone_scale is None:
-        return 1.0
-    if rig.girdle[j]:
-        return bone_scale[j]
-    canonical, rig_bind = rig.canonical_bone_length[j], rig.rig_bone_length[j]
-    if not (canonical > 1e-6) or not (rig_bind > 1e-6):
-        return bone_scale[j]
-    return bone_scale[j] * float(canonical / rig_bind)
-
-
-def studio_bone_worlds(globals_: torch.Tensor, posed_joints: torch.Tensor, rig: Rig,
-                       bone_scale=None, anchor_xz=None) -> torch.Tensor:
-    """Every exported bone's world matrix (T, B, 4, 4) in metres under Studio playback.
-
-    globals_: (T, 27, 3, 3) cskel27 global rotations; posed_joints: (T, 27, 3).
-    bone_scale: None reproduces a take without bone_scale (no arm stretch); a
-    27-vector reproduces a mocap take (grown offsets, stretched arm chain/hands).
-    anchor_xz: None (no anchoring) or (2,)/(T, 2) root XZ subtracted like the
-    anchor frame's root in applyMotionFrame.
-    """
-    T = globals_.shape[0]
-    dtype, device = globals_.dtype, globals_.device
-    bs = None if bone_scale is None else _bone_scale(bone_scale, globals_)
-    offsets = _studio_offsets(rig, bs)
-    s = rig.prep_scale
-    posed = posed_joints
-    if anchor_xz is not None:
-        a = torch.as_tensor(anchor_xz, dtype=dtype, device=device).reshape(-1, 2)
-        zero = torch.zeros_like(a[:, :1])
-        posed = posed - torch.cat((a[:, :1], zero, a[:, 1:]), dim=-1).unsqueeze(1)
-
-    desired = [None] * J
-    assumed_parent = {}
-    for j in range(J):
-        b = rig.prep_bone[j]
-        if b < 0:
-            continue
-        R = globals_[:, j]
-        pos = s * posed[:, j] + (R @ offsets[j].unsqueeze(-1)).squeeze(-1)
-        cp = rig.prep_chain_parent[j]
-        parent = desired[cp] @ rig.prep_chain_rel[j] if cp >= 0 else rig.prep_parent_bind_world[j].expand(T, 4, 4)
-        if rig.hierarchy_preserved[j]:
-            pos = _apply(parent, (rig.prep_bind_local_pos[j] * _stretch(rig, bs, j)).expand(T, 3))
-        rot = R @ rig.prep_bind_rot[j] * rig.prep_bind_scale[j]  # compose(pos, quat, scale)
-        desired[j] = _affine(rot, pos)
-        assumed_parent[b] = (parent, j)
-
-    leaves = {b: (joint, p, Rq) for b, joint, p, Rq in rig.leaves}
-    worlds = []
-    for b, parent_index in enumerate(rig.bone_parent):
-        if b in assumed_parent:
-            parent, j = assumed_parent[b]
-            local = torch.linalg.inv(parent) @ desired[j]
-        elif b in leaves:
-            joint, p, Rq = leaves[b]
-            local = _affine(Rq * rig.bone_rest_scale[b], p * _stretch(rig, bs, joint)).expand(T, 4, 4)
-        else:
-            local = rig.bone_rest_local[b].expand(T, 4, 4)
-        worlds.append(local if parent_index < 0 else worlds[parent_index] @ local)
-    W = torch.stack(worlds, dim=1)
-    root = torch.diag(torch.tensor([rig.root_scale] * 3 + [1.0], dtype=dtype, device=device))
-    return root @ W
-
-
-def lbs(bone_worlds: torch.Tensor, rig: Rig, rest: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
-    """Linear-blend skinning: v = sum_b w_vb * (W_b @ bindInverse_b) @ v_rest.
-
-    bone_worlds (T, B, 4, 4) metres; rest (V, 3); weights dense (V, B). -> (T, V, 3).
-    """
-    skin = bone_worlds @ rig.bind_inverse  # (T, B, 4, 4)
-    blended = torch.einsum("vb,tbij->tvij", weights, skin[:, :, :3, :])  # (T, V, 3, 4)
-    return (blended[..., :3] @ rest.unsqueeze(-1)).squeeze(-1) + blended[..., 3]
-
-
-def dense_weights(rig: Rig, indices=None) -> tuple:
-    """(rest (V, 3), dense weights (V, B)) of full-resolution vertices `indices` (all when None)."""
-    idx = slice(None) if indices is None else torch.as_tensor(indices, device=rig.vertices.device, dtype=torch.long)
-    si, sw = rig.skin_index[idx], rig.skin_weight[idx]
-    dense = torch.zeros(si.shape[0], len(rig.bone_names), dtype=sw.dtype, device=sw.device)
-    dense.scatter_add_(1, si, sw)
-    return rig.vertices[idx], dense
-
-
-def studio_skin(globals_: torch.Tensor, posed_joints: torch.Tensor, rig: Rig, bone_scale=None,
-                anchor_xz=None, vertex_indices=None):
-    """Layer 2. Returns (bone_worlds (T, B, 4, 4), vertices (T, V, 3)), metres.
-
-    Vertices are the decimated tracker mesh by default, or the full-resolution
-    export vertices `vertex_indices` (parity checks against Studio).
-    """
-    W = studio_bone_worlds(globals_, posed_joints, rig, bone_scale, anchor_xz)
-    if vertex_indices is None:
-        rest, weights = rig.dec_vertices, rig.dec_weights
-    else:
-        rest, weights = dense_weights(rig, vertex_indices)
-    return W, lbs(W, rig, rest, weights)