loop_window: the lowest-cost seam of a generated take, pose and velocity matched

Numpy only; the resident worker is its one caller. Hips-relative, heading-normalised
poses; cost is pose RMS plus 0.1 s of velocity RMS; ties go to the longer window,
then the earlier. Tests: a periodic stride between rests, the stated cost against a
frame-by-frame oracle, heading, ties, minimum length.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

scripts/loop_window.py

@@ -0,0 +1,41 @@
+"""The frame window of a generated take that loops with the smallest seam. Called by the resident worker, nowhere else.
+
+A seam is the jump from frame b back to frame a. Its cost is the pose difference plus 0.1 s times the velocity
+difference (Kovar, Gleicher and Pighin 2002, motion-graph transitions), both as RMS over joints, with each frame
+taken relative to its own Hips and turned to face a fixed way, so where the take wanders and which way it faces
+do not count. Numpy only.
+"""
+import numpy as np
+
+VELOCITY_WEIGHT_SECONDS = 0.1
+
+
+def loop_window(joints, heading, fps, min_frames, hips=0):
+    """joints [F, J, 3] world positions, heading [F] radians about Y (Kimodo's hip-vector angle), the joint index of
+    Hips. Returns (first, last, pose_cost, velocity_cost): frame indices into the take, last inclusive, last - first
+    at least min_frames, the pair with the lowest cost. A tie goes to the longer window, then to the earlier first."""
+    joints = np.asarray(joints, dtype=np.float64)
+    heading = np.asarray(heading, dtype=np.float64)
+    frames = joints.shape[0]
+    if joints.ndim != 3 or joints.shape[2] != 3 or heading.shape != (frames,) or not 1 <= min_frames < frames:
+        raise ValueError("loop_window: shapes or min_frames out of range")
+    relative = joints - joints[:, hips:hips + 1]
+    cos, sin = np.cos(-heading), np.sin(-heading)
+    x, z = relative[..., 0], relative[..., 2]
+    pose = relative.copy()
+    pose[..., 0] = x * cos[:, None] + z * sin[:, None]
+    pose[..., 2] = -x * sin[:, None] + z * cos[:, None]
+    velocity = np.gradient(pose, axis=0) * fps  # central differences, one-sided at the two ends
+
+    firsts, lasts = np.triu_indices(frames, k=min_frames)
+    pose_cost = np.empty(len(firsts))
+    velocity_cost = np.empty(len(firsts))
+    for first in range(frames - min_frames):  # one row of seams at a time: memory stays O(frames * joints)
+        row = slice(*np.searchsorted(firsts, [first, first + 1]))
+        later = lasts[row]
+        pose_cost[row] = np.sqrt(((pose[later] - pose[first]) ** 2).sum(-1).mean(-1))
+        velocity_cost[row] = np.sqrt(((velocity[later] - velocity[first]) ** 2).sum(-1).mean(-1))
+    cost = pose_cost + VELOCITY_WEIGHT_SECONDS * velocity_cost
+    best = np.flatnonzero(cost == cost.min())
+    pick = best[np.lexsort((firsts[best], -(lasts[best] - firsts[best])))[0]]
+    return int(firsts[pick]), int(lasts[pick]), float(pose_cost[pick]), float(velocity_cost[pick])

scripts/test_loop_window.py

@@ -0,0 +1,136 @@
+"""python3 -m unittest scripts/test_loop_window.py   (numpy only)"""
+import math
+import pathlib
+import sys
+import unittest
+
+import numpy as np
+
+sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
+from loop_window import loop_window  # noqa: E402
+
+FPS = 30
+
+
+def turn(points, angle):
+    """Rotate points [..., 3] about Y by angle, so that +Z turns toward +X (Kimodo's heading: 0 faces +Z)."""
+    x, z = points[..., 0], points[..., 2]
+    out = points.copy()
+    out[..., 0] = x * math.cos(angle) + z * math.sin(angle)
+    out[..., 2] = -x * math.sin(angle) + z * math.cos(angle)
+    return out
+
+
+def stride(frames=240, period=37, start=40, end=200):
+    """A stride cycle between two different rests. Hips travels; every other joint swings."""
+    rng = np.random.default_rng(3)
+    rest_a, rest_b = rng.normal(size=(2, 6, 3)) * 0.3
+    phase, direction = rng.uniform(0, 6.28, 6), rng.normal(size=(6, 3))
+    out = np.zeros((frames, 7, 3))
+    for f in range(frames):
+        weight = min(1.0, f / start, (frames - 1 - f) / (frames - 1 - end))
+        rest = rest_a if f < frames // 2 else rest_b
+        cycle = 0.2 * np.sin(2 * math.pi * f / period + phase)[:, None] * direction
+        out[f, 1:] = weight * (rest_a + cycle) + (1 - weight) * rest
+        out[f, 0] = (0.03 * f, 0.9, 0.0)  # Hips
+        out[f, 1:] += out[f, 0]
+    return out
+
+
+def reference(joints, heading, fps, min_frames):
+    """The cost written out frame by frame and joint by joint, to check the vectorised search against."""
+    frames, count = joints.shape[0], joints.shape[1]
+    pose = np.zeros((frames, count, 3))
+    for f in range(frames):
+        pose[f] = turn(joints[f] - joints[f, 0], -heading[f])
+    velocity = np.zeros_like(pose)
+    for f in range(frames):
+        lo, hi = max(f - 1, 0), min(f + 1, frames - 1)
+        velocity[f] = (pose[hi] - pose[lo]) / (hi - lo) * fps
+    best = None
+    for a in range(frames):
+        for b in range(a + min_frames, frames):
+            p = math.sqrt(sum(np.sum((pose[b, j] - pose[a, j]) ** 2) for j in range(count)) / count)
+            v = math.sqrt(sum(np.sum((velocity[b, j] - velocity[a, j]) ** 2) for j in range(count)) / count)
+            cost = p + 0.1 * v
+            key = (cost, -(b - a), a)
+            if best is None or key < best[0]:
+                best = (key, a, b, p, v)
+    return best[1:]
+
+
+def wandering(frames=44, count=4, seed=0):
+    """A smooth, unrepeating motion that turns as it goes, so the best window is unique and pose, velocity and
+    heading each move it."""
+    rng = np.random.default_rng(seed)
+    time = np.arange(frames)
+    joints = sum(rng.normal(size=(1, count, 3)) * 0.3 * np.sin(time[:, None, None] * rng.uniform(0.05, 0.5) + rng.uniform(0, 6.28))
+                 for _ in range(4))
+    heading = sum(0.4 * np.sin(time * rng.uniform(0.03, 0.2) + rng.uniform(0, 6.28)) for _ in range(2))
+    return joints, heading
+
+
+class LoopWindow(unittest.TestCase):
+    def test_a_periodic_motion_between_rest_ramps(self):
+        joints = stride()
+        first, last, pose_cost, _ = loop_window(joints, np.zeros(240), FPS, 120)
+        self.assertGreaterEqual(first, 40)
+        self.assertLessEqual(last, 200)
+        self.assertEqual((last - first) % 37, 0, (first, last))
+        self.assertGreaterEqual(last - first, 120)
+        self.assertLess(pose_cost, 1e-3)
+        self.assertNotEqual((first, last), (0, 239), "the whole take is not the answer")
+
+    def test_the_search_is_the_stated_cost(self):
+        joints, heading = wandering()
+        want_first, want_last, want_pose, want_velocity = reference(joints, heading, FPS, 22)
+        got = loop_window(joints, heading, FPS, 22)
+        self.assertEqual(got[:2], (want_first, want_last))
+        self.assertAlmostEqual(got[2], want_pose, places=9)
+        self.assertAlmostEqual(got[3], want_velocity, places=9)
+        # The velocity term decides here: matching poses alone would pick another window.
+        pose_only = min(
+            ((np.sqrt(((turn(joints[b] - joints[b, 0], -heading[b]) - turn(joints[a] - joints[a, 0], -heading[a])) ** 2).sum(-1).mean()), a, b)
+             for a in range(44) for b in range(a + 22, 44)))
+        self.assertNotEqual(pose_only[1:], (want_first, want_last))
+
+    def test_heading_is_normalised_away(self):
+        joints, heading = wandering()
+        upright = np.stack([turn(joints[f] - joints[f, :1], -heading[f]) + joints[f, :1] for f in range(44)])
+        turned = np.stack([turn(upright[f] - upright[f, :1], heading[f]) + upright[f, :1] for f in range(44)])
+        want = loop_window(upright, np.zeros(44), FPS, 22)
+        got = loop_window(turned, heading, FPS, 22)
+        self.assertEqual(got[:2], want[:2])
+        self.assertAlmostEqual(got[2], want[2], places=9)
+        self.assertNotEqual(loop_window(turned, np.zeros(44), FPS, 22)[:2], want[:2], "the turning would otherwise change the answer")
+
+    def test_a_tie_goes_to_the_longer_window_then_the_earlier(self):
+        cycle = np.array([[0, 0, 0], [0.5, 0.2, 0.1], [0.2, 0.7, 0.3], [0.4, 0.1, 0.9]], dtype=float)
+        joints = np.stack([np.stack([np.zeros(3), cycle[f % 4]]) for f in range(30)])
+        # Exact repeats every 4 frames: every a, b with b - a a multiple of 4 inside 1..28 costs zero.
+        self.assertEqual(loop_window(joints, np.zeros(30), FPS, 15)[:2], (1, 25))
+        still = np.zeros((30, 3, 3))
+        self.assertEqual(loop_window(still, np.zeros(30), FPS, 15)[:2], (0, 29))
+
+    def test_no_window_is_shorter_than_the_minimum(self):
+        joints, heading = wandering(frames=31)
+        for minimum in (5, 16, 25):
+            first, last, _, _ = loop_window(joints, heading, FPS, minimum)
+            self.assertGreaterEqual(last - first, minimum)
+        self.assertEqual(loop_window(joints, heading, FPS, 30)[:2], (0, 30))
+
+    def test_hips_is_the_joint_it_is_told(self):
+        joints, heading = wandering()
+        moved = joints[:, [2, 0, 1, 3]]
+        self.assertEqual(loop_window(moved, heading, FPS, 22, hips=1), loop_window(joints, heading, FPS, 22, hips=0))
+
+    def test_shapes_and_minimum_are_checked_without_echoing_them(self):
+        joints, heading = wandering()
+        for bad in ((joints, heading[:-1], 22), (joints, heading, 0), (joints, heading, 4321), (joints[..., :2], heading, 22)):
+            with self.assertRaises(ValueError) as caught:
+                loop_window(bad[0], bad[1], FPS, bad[2])
+            self.assertNotIn("4321", str(caught.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()