diff --git a/scripts/verify_ik_selfconsistency.py b/scripts/verify_ik_selfconsistency.py new file mode 100644 index 00000000..a6a56c7d --- /dev/null +++ b/scripts/verify_ik_selfconsistency.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python +"""Ground-truth self-consistency check for the inverse-kinematics pipeline. + +Each pre-extracted atomic batch stores BOTH the IK *input* (``keypoint_pos`` = +per-keypoint [pixel_x, pixel_y, depth_mm]) AND the *answer* that generated it +(``dof_angles`` = the simulated joint angles, named via the dataset's ``keys`` +attribute). That makes the IK directly testable: unproject the stored 2.5D +keypoints to 3D world coordinates, run the repo's IK, and check that the +recovered joint angles match the stored ground-truth angles. + +This validates, end to end: + * issue #48 / I3-A - the restored ``dof_name_lookup_canonical_to_nmf`` lets + ``_save_seqikpy_output`` run and pack DOFs correctly; + * issue #48 / I1-A - the camera unprojection uses the correct sensor size + (the keypoint labels are in ``rendering_size`` px); + * the IK solve itself (per-DOF angle error). + +It also has a flygym-free ``--emit-bounds`` mode that regenerates the +data-derived ``nmf_bounds`` (issue #48 / I3-B) from the observed ``dof_angles`` +range of motion. + +Usage +----- + # full self-consistency check (needs flygym + seqikpy installed): + python scripts/verify_ik_selfconsistency.py \ + --data-dir bulk_data/pose_estimation/atomic_batches/4variants \ + --n-batches 5 --max-frames 32 + + # regenerate data-derived joint bounds (no flygym needed): + python scripts/verify_ik_selfconsistency.py \ + --data-dir bulk_data/pose_estimation/atomic_batches/4variants \ + --emit-bounds --n-batches 500 + +NOTE (implemented by Claude Code (Opus 4.8) upon user instruction; issue #48). +""" + +from __future__ import annotations + +import argparse +import glob +from pathlib import Path + +import numpy as np + +# Camera unprojection only needs numpy/scipy, so it's safe to import eagerly. +from poseforge.pose.camera import CameraToWorldMapper + +# Production camera (see production/spotlight/keypoints3d.py and +# run_keypoints3d_inference.py). The atomic-batch ``keypoint_pos`` x,y are in the +# stored frame's pixel space (256 px for the shipped 4variants data), so the +# mapper MUST be built at that size -- this is exactly the I1-A fix. +DEFAULT_CAMERA = dict( + camera_pos=(0.0, 0.0, -100.0), + camera_fov_deg=5.0, + rotation_euler=(0.0, np.pi, -np.pi / 2.0), +) + +# Atomic-batch keypoint part names (NMF-style) -> canonical leg-keypoint names +# that ``invkin.run_seqikpy`` / ``keypoint_segments_canonical`` expect. +PART_TO_CANONICAL = { + "Coxa": "ThC", + "Femur": "CTr", + "Tibia": "FTi", + "Tarsus1": "TiTa", + "Tarsus5": "Claw", +} + +# The 7 canonical DOFs per leg, in the order the IK output is packed. +CANONICAL_DOFS = [ + "ThC_yaw", "ThC_pitch", "ThC_roll", + "CTr_pitch", "CTr_roll", "FTi_pitch", "TiTa_pitch", +] +LEGS = [f"{side}{pos}" for side in "LR" for pos in "FMH"] + + +def _decode_keys(attr) -> list[str]: + return [k.decode() if isinstance(k, bytes) else str(k) for k in attr] + + +def keypoint_keys_to_canonical(data_keys: list[str]) -> list[str]: + """Map atomic-batch keypoint names (e.g. ``LFCoxa``, ``LPedicel``) to the + canonical names IK expects (e.g. ``LFThC``, ``LPedicel``).""" + out = [] + for k in data_keys: + if k.endswith("Pedicel"): # antennae pass through unchanged + out.append(k) + continue + leg, part = k[:2], k[2:] + if part not in PART_TO_CANONICAL: + raise ValueError(f"Unrecognized keypoint part {part!r} in {k!r}") + out.append(f"{leg}{PART_TO_CANONICAL[part]}") + return out + + +def circular_abs_diff_deg(a_rad: np.ndarray, b_rad: np.ndarray) -> np.ndarray: + """|a - b| wrapped to [0, pi], in degrees (robust to angle wrapping).""" + d = (a_rad - b_rad + np.pi) % (2 * np.pi) - np.pi + return np.degrees(np.abs(d)) + + +def load_batches(data_dir: Path, n_batches: int, max_frames: int | None, seed: int): + files = sorted(glob.glob(str(data_dir / "**" / "*_labels.h5"), recursive=True)) + if not files: + raise FileNotFoundError(f"No *_labels.h5 under {data_dir}") + rng = np.random.default_rng(seed) + pick = files if n_batches >= len(files) else [ + files[i] for i in rng.choice(len(files), n_batches, replace=False) + ] + import h5py + + kp_list, dof_list, kp_keys, dof_keys = [], [], None, None + for f in pick: + with h5py.File(f, "r") as h: + kp = h["keypoint_pos"][:] + dof = h["dof_angles"][:] + if max_frames is not None: + kp, dof = kp[:max_frames], dof[:max_frames] + kp_list.append(kp) + dof_list.append(dof) + if kp_keys is None: + kp_keys = _decode_keys(h["keypoint_pos"].attrs["keys"]) + dof_keys = _decode_keys(h["dof_angles"].attrs["keys"]) + return np.concatenate(kp_list, 0), np.concatenate(dof_list, 0), kp_keys, dof_keys + + +def emit_bounds(data_dir: Path, n_batches: int, margin_deg: float, seed: int) -> None: + """Print a data-derived ``nmf_bounds`` dict (flygym-free). I3-B.""" + _, dof, _, dof_keys = load_batches(data_dir, n_batches, None, seed) + A = np.degrees(dof) + omin = {dof_keys[j]: A[:, j].min() for j in range(len(dof_keys))} + omax = {dof_keys[j]: A[:, j].max() for j in range(len(dof_keys))} + print(f"# data-derived from {A.shape[0]} frames, margin +/-{margin_deg:.0f} deg") + print("nmf_bounds = {") + for grp, pos in (("Front", "F"), ("Mid", "M"), ("Hind", "H")): + print(f" # {grp} legs") + for leg in [f"{s}{pos}" for s in "RL"]: + for dofn in CANONICAL_DOFS: + col = f"{leg}{dofn}" + if col not in omin: + continue + lo = np.floor(omin[col] - margin_deg) + hi = np.ceil(omax[col] + margin_deg) + wrap = " # WRAPPING: source RoM exceeds +/-180 deg" if ( + omin[col] < -180 or omax[col] > 180) else "" + print(f' "{leg}_{dofn}": (np.deg2rad({lo:.0f}), np.deg2rad({hi:.0f})),{wrap}') + print("}") + + +def run_selfconsistency(args) -> int: + kp, gt_dof, kp_keys, dof_keys = load_batches( + args.data_dir, args.n_batches, args.max_frames, args.seed + ) + n_frames = kp.shape[0] + print(f"Loaded {n_frames} frames; keypoints={len(kp_keys)} dofs={len(dof_keys)}") + + # 1) Unproject [px, py, depth] -> world mm with the CORRECT sensor size. + mapper = CameraToWorldMapper( + rendering_size=(args.rendering_size, args.rendering_size), **DEFAULT_CAMERA + ) + world = mapper(kp[..., :2], kp[..., 2]) # (n_frames, n_kpts, 3) mm + + # 2) Run the repo's IK. Imported lazily so --emit-bounds works without flygym. + import poseforge.pose.keypoints3d.invkin as invkin + + canonical_names = keypoint_keys_to_canonical(kp_keys) + joint_angles, _fk = invkin.run_seqikpy( + world_xyz=world, + keypoint_names_canonical=canonical_names, + n_workers=args.n_workers, + ) + + # 3) Compare recovered angles to ground truth, matching BY NAME (the IK output + # DOF order [yaw,pitch,roll,...] differs from the dof_angles order + # [pitch,roll,yaw,...], so positional comparison would be wrong). + gt_index = {k: j for j, k in enumerate(dof_keys)} + rows, all_err = [], [] + for leg in LEGS: + for dofn in CANONICAL_DOFS: + seqikpy_key = f"Angle_{leg}_{dofn}" + gt_key = f"{leg}{dofn}" + if seqikpy_key not in joint_angles or gt_key not in gt_index: + print(f" [skip] missing {seqikpy_key} or {gt_key}") + continue + rec = np.asarray(joint_angles[seqikpy_key]).reshape(-1)[:n_frames] + gt = gt_dof[:, gt_index[gt_key]] + err = circular_abs_diff_deg(rec, gt) + rows.append((gt_key, float(err.mean()), float(np.percentile(err, 90)))) + all_err.append(err) + + all_err = np.concatenate(all_err) + print(f"\n{'DOF':16s} {'MAE(deg)':>9s} {'p90(deg)':>9s}") + for name, mae, p90 in sorted(rows, key=lambda r: -r[1]): + print(f"{name:16s} {mae:9.2f} {p90:9.2f}") + print(f"\nOVERALL mean abs angle error: {all_err.mean():.2f} deg " + f"(median {np.median(all_err):.2f}, p90 {np.percentile(all_err, 90):.2f})") + print(f"fraction of DOF-frames within 5 deg: {(all_err < 5).mean() * 100:.1f}%") + # Heuristic pass/fail: a correct mapping + camera + IK should sit well under + # ~15 deg mean; gross mapping/ordering/camera bugs blow this up. + threshold = args.threshold_deg + ok = all_err.mean() < threshold + print(f"\n{'PASS' if ok else 'FAIL'}: overall mean < {threshold:.0f} deg") + return 0 if ok else 1 + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--data-dir", type=Path, + default=Path("bulk_data/pose_estimation/atomic_batches/4variants")) + p.add_argument("--n-batches", type=int, default=5, + help="number of atomic batches to sample") + p.add_argument("--max-frames", type=int, default=32, + help="max frames per batch (None=all)") + p.add_argument("--rendering-size", type=int, default=256, + help="pixel size of the stored keypoint labels (sensor size)") + p.add_argument("--n-workers", type=int, default=6) + p.add_argument("--threshold-deg", type=float, default=15.0, + help="pass if overall mean angle error is below this") + p.add_argument("--seed", type=int, default=0) + p.add_argument("--emit-bounds", action="store_true", + help="print data-derived nmf_bounds and exit (no flygym needed)") + p.add_argument("--margin-deg", type=float, default=10.0, + help="padding added to observed RoM when emitting bounds") + args = p.parse_args() + + if args.emit_bounds: + emit_bounds(args.data_dir, args.n_batches, args.margin_deg, args.seed) + return 0 + return run_selfconsistency(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/poseforge/neuromechfly/constants.py b/src/poseforge/neuromechfly/constants.py index 76df9a21..327db0a1 100644 --- a/src/poseforge/neuromechfly/constants.py +++ b/src/poseforge/neuromechfly/constants.py @@ -18,6 +18,47 @@ keypoint_name_lookup_canonical_to_nmf = { v: k for k, v in keypoint_name_lookup_nmf_to_canonical.items() } + +# Mapping from the canonical (Aymanns et al. 2022) leg DOF names to the +# NeuroMechFly DOF (joint) names. +# +# This constant was originally introduced in commit 5f8b6ee and accidentally +# removed in 3c9449a ("compatibility with flygym v2") while its call sites in +# `run_inverse_kinematics.py` and `production/spotlight/keypoints3d.py` were +# never updated -> they raised AttributeError when saving IK output. Restored +# here (see issue #48, finding I3-A). +# +# IMPORTANT - ordering and key names: +# * The KEYS are the 7 canonical leg DOF names. They are kept in the order +# ThC_yaw, ThC_pitch, ThC_roll, CTr_pitch, CTr_roll, FTi_pitch, TiTa_pitch +# because the call sites do +# `dof_names_per_leg = list(dof_name_lookup_canonical_to_nmf.keys())` +# and use that order as the DOF axis of the saved `joint_angles` array +# (and as its `dof_names_per_leg` attribute). This is the same DOF ordering +# used by `nmf_initial_angles` / seqikpy's kinematic chain (yaw, pitch, +# roll, ...). +# * seqikpy emits joint-angle dict keys of the form +# `Angle_{leg}_{canonical_dof}` (e.g. "Angle_LF_ThC_yaw"); see +# seqikpy.leg_inverse_kinematics.LegInvKinSeq. The call sites build the +# lookup key as `f"Angle_{leg}_{dof_name}"` where `dof_name` is a KEY of +# this dict, so the keys must be exactly these canonical DOF names for the +# lookups (and hence the DOF packing order) to be correct. +# * The VALUES are the corresponding NeuroMechFly DOF names. They are not used +# by the IK save path today but document the canonical<->NMF correspondence +# and keep this dict useful for downstream NMF actuation code. +dof_name_lookup_canonical_to_nmf = { + "ThC_yaw": "Coxa_yaw", + "ThC_pitch": "Coxa", + "ThC_roll": "Coxa_roll", + "CTr_pitch": "Femur", + "CTr_roll": "Femur_roll", + "FTi_pitch": "Tibia", + "TiTa_pitch": "Tarsus1", +} +dof_name_lookup_nmf_to_canonical = { + v: k for k, v in dof_name_lookup_canonical_to_nmf.items() +} + legs = [f"{side}{pos}" for side in "LR" for pos in "FMH"] leg_keypoints_canonical = ["ThC", "CTr", "FTi", "TiTa", "Claw"] leg_keypoints_nmf = [keypoint_name_lookup_canonical_to_nmf[kp] for kp in leg_keypoints_canonical] @@ -284,55 +325,79 @@ def parse_nmf_joint(joint: JointDOF) -> tuple[str, str]: "LH_Claw": np.array([-0.215, 0.087, -2.588]), } -# Determine the bounds for each joint DOF +# Joint DOF bounds for seqikpy IK. DATA-DERIVED (issue #48 I3-B). +# +# Method (values are whole degrees wrapped in np.deg2rad): +# 1. Source: the ground-truth simulated `dof_angles` stored in every atomic-batch +# `_labels.h5` (named via the dataset's `keys` attr; 6 legs x 7 DOFs = 42). +# 2. Sample: 500 `_labels.h5` files (numpy default_rng(0), no replacement) from the +# sorted recursive glob of bulk_data/.../atomic_batches/4variants/**/*_labels.h5, +# all 32 frames each -> n = 16000 frames. +# 3. Per DOF: bound = (floor(min_deg - margin), ceil(max_deg + margin)) with +# margin = 10 deg -- the observed range padded outward. The margin gives the +# least-squares solver headroom so the optimum does not sit exactly on a boundary +# (an IK fragility noted in the audit); min/max (not percentiles) guarantee no +# observed pose is clipped. +# 4. Map ground-truth key `{leg}{dof}` (e.g. RFThC_yaw) -> bounds key `{leg}_{dof}`. +# Reproduce EXACTLY with: +# python scripts/verify_ik_selfconsistency.py --emit-bounds --n-batches 500 --seed 0 --margin-deg 10 +# +# Why: supersedes the earlier hand-set / L-R-mirrored bounds, 10/42 of which were tighter +# than the actual range of motion and would clip valid poses (a bound tighter than the data +# is unreachable by IK -> forces a wrong solution). The simulated dof_angles are exactly the +# RoM the IK must reproduce, so they are the correct floor for these bounds. +# CAVEATS: (1) this is the *training* RoM, not the anatomical RoM; production may see novel +# poses, so widen toward NeuroMechFly's anatomical limits if IK saturates a bound. +# (2) DOFs flagged WRAPPING below have source angles beyond +/-180 deg, indicating the +# upstream kinematics need unwrapping; bounds contain them only so IK can reproduce them. +# Wrapping DOFs: RH_ThC_roll. (3) Bounds come out near-mirror L/R where the data is, but +# are NOT forced symmetric (data-honest). nmf_bounds = { # Front legs - "RF_ThC_yaw": (np.deg2rad(-45), np.deg2rad(45)), - "RF_ThC_pitch": (np.deg2rad(-10), np.deg2rad(90)), - "RF_ThC_roll": (np.deg2rad(-135), np.deg2rad(10)), # ? 1 - "RF_CTr_pitch": (np.deg2rad(-270), np.deg2rad(10)), # ? 2 - "RF_CTr_roll": (np.deg2rad(-180), np.deg2rad(90)), # ? 3 - "RF_FTi_pitch": (np.deg2rad(-10), np.deg2rad(180)), - "RF_TiTa_pitch": (np.deg2rad(-180), np.deg2rad(10)), - "LF_ThC_yaw": (np.deg2rad(-45), np.deg2rad(45)), - "LF_ThC_pitch": (np.deg2rad(-10), np.deg2rad(90)), - "LF_ThC_roll": (np.deg2rad(-10), np.deg2rad(90)), # ? 1 - "LF_CTr_pitch": (np.deg2rad(-180), np.deg2rad(10)), # ? 2 - "LF_CTr_roll": (np.deg2rad(-90), np.deg2rad(180)), # ? 3 - "LF_FTi_pitch": (np.deg2rad(-10), np.deg2rad(180)), - "LF_TiTa_pitch": (np.deg2rad(-180), np.deg2rad(10)), - + "RF_ThC_yaw": (np.deg2rad(-36), np.deg2rad(57)), + "RF_ThC_pitch": (np.deg2rad(-17), np.deg2rad(79)), + "RF_ThC_roll": (np.deg2rad(-185), np.deg2rad(105)), + "RF_CTr_pitch": (np.deg2rad(-187), np.deg2rad(-41)), + "RF_CTr_roll": (np.deg2rad(-187), np.deg2rad(13)), + "RF_FTi_pitch": (np.deg2rad(5), np.deg2rad(178)), + "RF_TiTa_pitch": (np.deg2rad(-147), np.deg2rad(9)), + "LF_ThC_yaw": (np.deg2rad(-63), np.deg2rad(37)), + "LF_ThC_pitch": (np.deg2rad(-21), np.deg2rad(65)), + "LF_ThC_roll": (np.deg2rad(-18), np.deg2rad(172)), + "LF_CTr_pitch": (np.deg2rad(-180), np.deg2rad(-51)), + "LF_CTr_roll": (np.deg2rad(-21), np.deg2rad(186)), + "LF_FTi_pitch": (np.deg2rad(-6), np.deg2rad(182)), + "LF_TiTa_pitch": (np.deg2rad(-150), np.deg2rad(10)), # Mid legs - "RM_ThC_yaw": (np.deg2rad(-45), np.deg2rad(45)), # ? 4 - "RM_ThC_pitch": (np.deg2rad(-10), np.deg2rad(90)), - "RM_ThC_roll": (np.deg2rad(-180), np.deg2rad(10)), # ? 5 - "RM_CTr_pitch": (np.deg2rad(-270), np.deg2rad(10)), # ? 6 - "RM_CTr_roll": (np.deg2rad(-90), np.deg2rad(90)), - "RM_FTi_pitch": (np.deg2rad(-10), np.deg2rad(180)), - "RM_TiTa_pitch": (np.deg2rad(-180), np.deg2rad(10)), - "LM_ThC_yaw": (np.deg2rad(-45), np.deg2rad(90)), # ? 4 - "LM_ThC_pitch": (np.deg2rad(-10), np.deg2rad(90)), - "LM_ThC_roll": (np.deg2rad(-10), np.deg2rad(180)), # ? 5 - "LM_CTr_pitch": (np.deg2rad(-180), np.deg2rad(10)), # ? 6 - "LM_CTr_roll": (np.deg2rad(-90), np.deg2rad(90)), - "LM_FTi_pitch": (np.deg2rad(-10), np.deg2rad(180)), - "LM_TiTa_pitch": (np.deg2rad(-180), np.deg2rad(10)), - + "RM_ThC_yaw": (np.deg2rad(-45), np.deg2rad(28)), + "RM_ThC_pitch": (np.deg2rad(-31), np.deg2rad(27)), + "RM_ThC_roll": (np.deg2rad(-171), np.deg2rad(-29)), + "RM_CTr_pitch": (np.deg2rad(-155), np.deg2rad(-50)), + "RM_CTr_roll": (np.deg2rad(-78), np.deg2rad(12)), + "RM_FTi_pitch": (np.deg2rad(0), np.deg2rad(166)), + "RM_TiTa_pitch": (np.deg2rad(-81), np.deg2rad(9)), + "LM_ThC_yaw": (np.deg2rad(-26), np.deg2rad(38)), + "LM_ThC_pitch": (np.deg2rad(-32), np.deg2rad(26)), + "LM_ThC_roll": (np.deg2rad(35), np.deg2rad(167)), + "LM_CTr_pitch": (np.deg2rad(-157), np.deg2rad(-38)), + "LM_CTr_roll": (np.deg2rad(-13), np.deg2rad(83)), + "LM_FTi_pitch": (np.deg2rad(4), np.deg2rad(163)), + "LM_TiTa_pitch": (np.deg2rad(-132), np.deg2rad(10)), # Hind legs - "RH_ThC_yaw": (np.deg2rad(-45), np.deg2rad(45)), # ? 7 - "RH_ThC_pitch": (np.deg2rad(-10), np.deg2rad(90)), - "RH_ThC_roll": (np.deg2rad(-180), np.deg2rad(10)), # ? 8 - "RH_CTr_pitch": (np.deg2rad(-180), np.deg2rad(10)), - "RH_CTr_roll": (np.deg2rad(-90), np.deg2rad(90)), - "RH_FTi_pitch": (np.deg2rad(-10), np.deg2rad(180)), - "RH_TiTa_pitch": (np.deg2rad(-180), np.deg2rad(10)), - "LH_ThC_yaw": (np.deg2rad(-45), np.deg2rad(90)), # ? 7 - "LH_ThC_pitch": (np.deg2rad(-10), np.deg2rad(90)), - "LH_ThC_roll": (np.deg2rad(-10), np.deg2rad(180)), # ? 8 - "LH_CTr_pitch": (np.deg2rad(-180), np.deg2rad(10)), - "LH_CTr_roll": (np.deg2rad(-90), np.deg2rad(90)), - "LH_FTi_pitch": (np.deg2rad(-10), np.deg2rad(180)), - "LH_TiTa_pitch": (np.deg2rad(-180), np.deg2rad(10)), + "RH_ThC_yaw": (np.deg2rad(-65), np.deg2rad(39)), + "RH_ThC_pitch": (np.deg2rad(-33), np.deg2rad(50)), + "RH_ThC_roll": (np.deg2rad(-216), np.deg2rad(-62)), # WRAPPING: source RoM exceeds +/-180 deg + "RH_CTr_pitch": (np.deg2rad(-158), np.deg2rad(-27)), + "RH_CTr_roll": (np.deg2rad(-16), np.deg2rad(167)), + "RH_FTi_pitch": (np.deg2rad(-3), np.deg2rad(168)), + "RH_TiTa_pitch": (np.deg2rad(-156), np.deg2rad(10)), + "LH_ThC_yaw": (np.deg2rad(-28), np.deg2rad(66)), + "LH_ThC_pitch": (np.deg2rad(-33), np.deg2rad(53)), + "LH_ThC_roll": (np.deg2rad(63), np.deg2rad(187)), + "LH_CTr_pitch": (np.deg2rad(-169), np.deg2rad(-10)), + "LH_CTr_roll": (np.deg2rad(-115), np.deg2rad(85)), + "LH_FTi_pitch": (np.deg2rad(2), np.deg2rad(168)), + "LH_TiTa_pitch": (np.deg2rad(-123), np.deg2rad(9)), } nmf_size = { diff --git a/src/poseforge/pose/keypoints3d/invkin.py b/src/poseforge/pose/keypoints3d/invkin.py index 9355af1e..792d32b0 100644 --- a/src/poseforge/pose/keypoints3d/invkin.py +++ b/src/poseforge/pose/keypoints3d/invkin.py @@ -11,12 +11,74 @@ from poseforge.pose.keypoints3d.visualizer import visualize_leg_segment_lengths +def _interpolate_nan_frames( + data_block: np.ndarray, +) -> tuple[np.ndarray, int, int]: + """Linearly interpolate (and edge-fill) NaN values along the time axis. + + seqikpy's IK solver has no NaN handling: a single NaN/occluded keypoint that + reaches the solver propagates to NaN joint angles, which then trips the + ``assert not np.isnan(...)`` at IK save time and aborts the whole recording + (issue #48, finding I3-D). Rather than aborting, we fill short occlusion gaps + per keypoint/coordinate so the solver always receives finite input. + + The fill is per (keypoint, coordinate) time series: + * interior NaNs are linearly interpolated from the nearest finite + neighbours on each side, and + * leading/trailing NaNs are filled with the nearest finite value + (forward/backward fill at the edges). + + Args: + data_block: array of shape (n_frames, n_keypoints, 3). Modified on a copy. + + Returns: + (filled, n_nan_values, n_affected_frames): the filled array (a copy), the + number of NaN scalar values present before filling, and the number of + frames that had at least one NaN coordinate before filling. If an entire + time series for a (keypoint, coordinate) is NaN it cannot be filled and + remains NaN; the caller is responsible for surfacing that. + """ + filled = data_block.astype(np.float32, copy=True) + n_frames = filled.shape[0] + nan_mask = np.isnan(filled) + n_nan_values = int(nan_mask.sum()) + n_affected_frames = int(nan_mask.any(axis=(1, 2)).sum()) + if n_nan_values == 0: + return filled, 0, 0 + + frame_idx = np.arange(n_frames) + # Flatten (keypoint, coord) into independent time series. + flat = filled.reshape(n_frames, -1) + flat_nan = nan_mask.reshape(n_frames, -1) + for series_idx in range(flat.shape[1]): + col_nan = flat_nan[:, series_idx] + if not col_nan.any(): + continue + valid = ~col_nan + if not valid.any(): + # Entire series is NaN; nothing to interpolate from. Leave as NaN. + continue + # np.interp clamps to the first/last valid value at the edges, which + # gives us forward/backward fill for leading/trailing NaNs for free. + flat[col_nan, series_idx] = np.interp( + frame_idx[col_nan], frame_idx[valid], flat[valid, series_idx] + ) + return filled, n_nan_values, n_affected_frames + + def _world_xyz_to_seqikpy_format( world_xyz: np.ndarray, keypoint_names_canonical: list[str] | np.ndarray, max_n_frames: int | None = None, ) -> dict[str, np.ndarray]: - """Convert raw 3D keypoint positions to format expected by SeqIKPy.""" + """Convert raw 3D keypoint positions to format expected by SeqIKPy. + + Occluded/NaN keypoints are not fatal: NaN gaps are interpolated per leg over + time and a warning is logged with the count and location (issue #48, finding + I3-D). Only a keypoint that is NaN for *every* frame of a recording is left + as NaN and aborts (it cannot be recovered and would otherwise silently corrupt + the whole leg chain). + """ n_frames, n_keypoints, _ = world_xyz.shape if max_n_frames is not None: n_frames = min(n_frames, max_n_frames) @@ -34,7 +96,30 @@ def _world_xyz_to_seqikpy_format( poseforge_key = f"{leg}{nmf_constants.keypoint_name_lookup_nmf_to_canonical[keypoint_name]}" idx = keypoint_names_canonical.index(poseforge_key) data_block[:, keypoint_idx, :] = world_xyz[:n_frames, idx, :] - assert not np.isnan(data_block).any() + + # Gracefully handle NaN/occluded keypoints instead of aborting the whole + # recording on a single NaN (was: `assert not np.isnan(data_block).any()`). + data_block, n_nan_values, n_affected_frames = _interpolate_nan_frames( + data_block + ) + if n_nan_values > 0: + logger.warning( + f"Leg {leg}: found {n_nan_values} NaN coordinate value(s) " + f"(occluded keypoints) in {n_affected_frames}/{n_frames} frames; " + f"linearly interpolated over time before inverse kinematics." + ) + # Any remaining NaN means a keypoint was occluded for the entire recording + # and could not be recovered. This would corrupt the IK chain, so fail loud + # and clear (per-leg) rather than producing silently wrong angles. + if np.isnan(data_block).any(): + fully_nan = np.isnan(data_block).all(axis=0) # (n_keypoints, 3) + bad_kp_idxs = sorted({int(i) for i, _ in zip(*np.where(fully_nan))}) + bad_kps = [nmf_constants.leg_keypoints_nmf[i] for i in bad_kp_idxs] + raise ValueError( + f"Leg {leg}: keypoint(s) {bad_kps} are NaN for the entire " + f"recording and cannot be interpolated. Cannot run inverse " + f"kinematics for this leg." + ) pose_data_dict[f"{leg}_leg"] = data_block return pose_data_dict @@ -186,12 +271,86 @@ def run_seqikpy( return joint_angles, forward_kinematics +def detect_large_joint_angle_jumps( + joint_angles: dict[str, np.ndarray], + threshold_rad: float = np.deg2rad(45.0), +) -> dict[str, np.ndarray]: + """Detect large frame-to-frame jumps in per-DOF joint-angle time series. + + seqikpy seeds the IK solve at frame ``t`` with the solution from frame + ``t-1``, so a single bad solve can propagate; and with + ``parallel_over_time=True`` each time chunk is re-seeded from the static + initial angles and chunks are linearly blended, which can introduce wrong + transients at chunk boundaries (issue #48, finding I3-C). + + This is a *post-hoc* detector: it does not fix the angles, it flags frames + whose absolute change from the previous frame exceeds ``threshold_rad`` so a + caller can log / inspect them. It is intentionally pure (dict of numpy arrays + in, dict of boolean masks out) so it is easy to test and reuse. + + Args: + joint_angles: mapping of ``Angle_{leg}_{dof}`` -> 1D array of radians. + threshold_rad: absolute per-frame jump (radians) above which a frame is + flagged. Defaults to 45 degrees. + + Returns: + Mapping from each input key to a boolean mask of shape (n_frames,) that is + ``True`` at frame ``t`` when ``|angle[t] - angle[t-1]| > threshold_rad``. + Frame 0 is always ``False`` (no previous frame). Keys whose values are not + 1D arrays are skipped. + """ + jumps: dict[str, np.ndarray] = {} + for key, series in joint_angles.items(): + series = np.asarray(series) + if series.ndim != 1 or series.shape[0] < 2: + continue + diff = np.abs(np.diff(series)) + mask = np.zeros(series.shape[0], dtype=bool) + mask[1:] = diff > threshold_rad + jumps[key] = mask + return jumps + + +def log_large_joint_angle_jumps( + joint_angles: dict[str, np.ndarray], + threshold_rad: float = np.deg2rad(45.0), +) -> int: + """Run :func:`detect_large_joint_angle_jumps` and log a warning summary. + + Returns the total number of flagged (DOF, frame) pairs. Logs nothing beyond a + debug line when no jumps are found. + """ + jumps = detect_large_joint_angle_jumps(joint_angles, threshold_rad=threshold_rad) + total = int(sum(int(mask.sum()) for mask in jumps.values())) + if total == 0: + logger.debug( + f"No frame-to-frame joint-angle jumps above " + f"{np.rad2deg(threshold_rad):.0f} deg detected." + ) + return 0 + + # Summarise per DOF (only those with at least one jump), most jumps first. + per_dof = {k: int(m.sum()) for k, m in jumps.items() if m.any()} + summary = ", ".join( + f"{k}:{n}" for k, n in sorted(per_dof.items(), key=lambda kv: -kv[1]) + ) + logger.warning( + f"Detected {total} large frame-to-frame joint-angle jump(s) " + f"(>{np.rad2deg(threshold_rad):.0f} deg). These can indicate a bad IK " + f"solve propagating or a chunk-boundary transient (parallel_over_time). " + f"Per-DOF counts: {summary}. For correctness-critical runs, consider " + f"re-running with parallel_over_time=False (or n_workers=1)." + ) + return total + + def align_fwdkin_xyz_to_rawpred_xyz( keypoints_pos_raw: np.ndarray, keypoints_pos_constrained: np.ndarray, keypoints_order: list[str], legs: list[str], leg_keypoints_canonical: list[str], + keypoints_order_constrained: list[str] | None = None, ) -> np.ndarray: """Align constrained poses to raw poses by shifting each leg's kinematic chain. @@ -199,52 +358,84 @@ def align_fwdkin_xyz_to_rawpred_xyz( we want to shift each leg back so that the first keypoint (ThC/Coxa) has the same 3D position as in the raw poses. + The raw and constrained arrays may use *different* keypoint orderings: the raw + array is ordered by the inference HDF5's ``keypoints`` attribute, while the + constrained array is produced by :func:`fwdkin_world_xyz_append_antennae`, which + returns its own ordering. Previously this function indexed *both* arrays with the + raw ``keypoints_order``, which silently produced wrong results if the two orders + ever differed (issue #48, finding I3-E). We now index each array with its own + order and validate the constrained order. + Args: keypoints_pos_raw: Raw keypoint positions (n_frames, n_keypoints, 3) keypoints_pos_constrained: Constrained keypoint positions (n_frames, n_keypoints, 3) - keypoints_order: List of keypoint names + keypoints_order: List of keypoint names for ``keypoints_pos_raw`` legs: List of leg names ['LF', 'LM', 'LH', 'RF', 'RM', 'RH'] leg_keypoints_canonical: List of keypoint names per leg ['ThC', 'CTr', 'FTi', 'TiTa', 'Claw'] + keypoints_order_constrained: List of keypoint names for + ``keypoints_pos_constrained``. Defaults to ``keypoints_order`` for + backwards compatibility (i.e. the caller asserts both arrays share an + ordering). Returns: keypoints_pos_constrained_aligned: Aligned constrained poses """ + if keypoints_order_constrained is None: + keypoints_order_constrained = keypoints_order + keypoints_pos_constrained_aligned = keypoints_pos_constrained.copy() n_frames = keypoints_pos_raw.shape[0] # For each leg, align the constrained pose to the raw pose for leg in legs: - # Get the first keypoint (ThC/Coxa) for this leg + # Get the first keypoint (ThC/Coxa) for this leg, looked up separately in + # each array's own keypoint ordering. first_keypoint_name = f"{leg}{leg_keypoints_canonical[0]}" # e.g., "LFThC" try: - first_keypoint_idx = keypoints_order.index(first_keypoint_name) + first_keypoint_idx_raw = keypoints_order.index(first_keypoint_name) except ValueError: logger.warning( - f"Keypoint {first_keypoint_name} not found in keypoints_order" + f"Keypoint {first_keypoint_name} not found in raw keypoints_order" + ) + continue + try: + first_keypoint_idx_constr = keypoints_order_constrained.index( + first_keypoint_name + ) + except ValueError: + logger.warning( + f"Keypoint {first_keypoint_name} not found in " + f"keypoints_order_constrained" ) continue - # Get all keypoint indices for this leg - leg_keypoint_indices = [] + # Get all keypoint indices for this leg, in both orderings. We keep them + # paired so the translation applies to the same physical keypoint in each + # array even if the orderings differ. + leg_keypoint_index_pairs: list[tuple[int, int]] = [] for keypoint in leg_keypoints_canonical: keypoint_name = f"{leg}{keypoint}" try: - idx = keypoints_order.index(keypoint_name) - leg_keypoint_indices.append(idx) + idx_raw = keypoints_order.index(keypoint_name) + idx_constr = keypoints_order_constrained.index(keypoint_name) except ValueError: - logger.warning(f"Keypoint {keypoint_name} not found in keypoints_order") + logger.warning( + f"Keypoint {keypoint_name} not found in raw/constrained " + f"keypoints_order" + ) continue + leg_keypoint_index_pairs.append((idx_raw, idx_constr)) - if not leg_keypoint_indices: + if not leg_keypoint_index_pairs: continue # For each frame, compute the translation needed to align the first keypoint for frame_idx in range(n_frames): # Get the positions of the first keypoint in raw and constrained poses - raw_first_pos = keypoints_pos_raw[frame_idx, first_keypoint_idx] + raw_first_pos = keypoints_pos_raw[frame_idx, first_keypoint_idx_raw] constrained_first_pos = keypoints_pos_constrained[ - frame_idx, first_keypoint_idx + frame_idx, first_keypoint_idx_constr ] # Skip if either position has NaN values @@ -254,11 +445,14 @@ def align_fwdkin_xyz_to_rawpred_xyz( # Compute translation vector translation = raw_first_pos - constrained_first_pos - # Apply translation to all keypoints of this leg - for leg_kp_idx in leg_keypoint_indices: - current_pos = keypoints_pos_constrained_aligned[frame_idx, leg_kp_idx] + # Apply translation to all keypoints of this leg (indexed in the + # constrained array's own ordering). + for _, leg_kp_idx_constr in leg_keypoint_index_pairs: + current_pos = keypoints_pos_constrained_aligned[ + frame_idx, leg_kp_idx_constr + ] if not np.isnan(current_pos).any(): - keypoints_pos_constrained_aligned[frame_idx, leg_kp_idx] = ( + keypoints_pos_constrained_aligned[frame_idx, leg_kp_idx_constr] = ( current_pos + translation ) diff --git a/src/poseforge/pose/keypoints3d/scripts/run_inverse_kinematics.py b/src/poseforge/pose/keypoints3d/scripts/run_inverse_kinematics.py index 79fccd81..64f4c501 100644 --- a/src/poseforge/pose/keypoints3d/scripts/run_inverse_kinematics.py +++ b/src/poseforge/pose/keypoints3d/scripts/run_inverse_kinematics.py @@ -523,6 +523,10 @@ def visualize_inverse_kinematics_comparison( keypoints_order=keypoints_order, legs=legs_ik, leg_keypoints_canonical=leg_keypoints_canonical_ik, + # The constrained array is built by fwdkin_world_xyz_append_antennae and + # has its own keypoint ordering; pass it explicitly so each array is + # indexed by its own order rather than assuming they match (I3-E, #48). + keypoints_order_constrained=keypoints_order_constrained, ) # For now, we don't have 2D projections of the constrained poses, so set to None @@ -594,7 +598,37 @@ def process_all( n_workers_per_dataset: int = 6, create_visualization: bool = False, input_images_basedir: str | None = None, + correctness_critical: bool = False, + jump_warning_threshold_deg: float = 45.0, + seqikpy_kwargs: dict | None = None, ) -> None: + """Run inverse kinematics on all keypoints3d outputs under ``input_dirs``. + + Args: + correctness_critical: If True, run seqikpy without time-chunk + parallelization/blending (``parallel_over_time=False``). seqikpy seeds + frame ``t`` from frame ``t-1``; with ``parallel_over_time=True`` time + chunks are re-seeded from static initial angles and linearly blended, + which can produce wrong transients at chunk boundaries (issue #48, + finding I3-C). Leg-level parallelism (``n_workers_per_dataset``) is + still used, so this only forgoes time-axis chunking. Ignored for keys + already present in ``seqikpy_kwargs``. + jump_warning_threshold_deg: Frames whose joint angle changes by more than + this (degrees) from the previous frame are flagged and logged after + each IK solve (post-hoc I3-C detection; does not modify the angles). + seqikpy_kwargs: Extra keyword arguments forwarded to + ``LegInvKinSeq.run_ik_and_fk`` (e.g. ``parallel_over_time``, + ``chunk_overlap``, ``min_chunk_size``). Takes precedence over the + flags above. + """ + # Assemble the kwargs forwarded to seqikpy's run_ik_and_fk. + run_ik_and_fk_kwargs: dict = {} + if correctness_critical: + # Disable time-chunk parallelization + blending (still parallel over legs). + run_ik_and_fk_kwargs["parallel_over_time"] = False + if seqikpy_kwargs: + run_ik_and_fk_kwargs.update(seqikpy_kwargs) + # Index all keypoints3d output files to process all_keypoints3d_output_files = [] for input_dir in input_dirs: @@ -615,6 +649,14 @@ def process_all( max_n_frames=max_n_frames, n_workers=n_workers_per_dataset, debug_plots_dir=keypoints3d_output_file.parent / "ik_debug_plots/", + **run_ik_and_fk_kwargs, + ) + # Post-hoc detection + logging of large frame-to-frame joint-angle jumps + # (I3-C). This flags likely bad-solve propagation / chunk-boundary + # transients; it does not alter the saved angles. + invkin.log_large_joint_angle_jumps( + joint_angles, + threshold_rad=np.deg2rad(jump_warning_threshold_deg), ) _save_seqikpy_output( output_path, joint_angles, forward_kinematics, frame_ids=frame_ids diff --git a/tests/test_ik_robustness.py b/tests/test_ik_robustness.py new file mode 100644 index 00000000..9bb1d4cd --- /dev/null +++ b/tests/test_ik_robustness.py @@ -0,0 +1,359 @@ +"""Pure-logic tests for the inverse-kinematics robustness fixes (issue #48, +findings I3-A .. I3-E). + +These tests deliberately avoid importing ``poseforge.neuromechfly.constants`` or +``poseforge.pose.keypoints3d.invkin`` directly: those modules import ``flygym`` +(``from flygym.anatomy import JointDOF``), which is NOT installed in this +environment (it does not support Python 3.13 yet). Importing them would fail at +collection time. + +Instead we: + * parse the relevant source files with ``ast`` and assert on the literal + definitions (the restored DOF map I3-A, and the mirror-consistency of + ``nmf_bounds`` I3-B), and + * extract the *actual source* of the pure numpy helpers added for I3-C / I3-D + and ``exec`` them in an isolated namespace (numpy + a stub logger) so the + tests exercise the real implementation rather than a replica. + +NOTE (testing limitation): full IK import/run testing is blocked by the missing +``flygym`` dependency, so these tests cover the pure logic and the static +definitions only. +""" + +import ast +import textwrap +from pathlib import Path + +import numpy as np +import pytest + + +# --------------------------------------------------------------------------- # +# Locate source files relative to this test (repo-root/src/poseforge/...). +# --------------------------------------------------------------------------- # +REPO_ROOT = Path(__file__).resolve().parents[1] +CONSTANTS_PY = REPO_ROOT / "src" / "poseforge" / "neuromechfly" / "constants.py" +INVKIN_PY = REPO_ROOT / "src" / "poseforge" / "pose" / "keypoints3d" / "invkin.py" +RUN_IK_PY = ( + REPO_ROOT + / "src" + / "poseforge" + / "pose" + / "keypoints3d" + / "scripts" + / "run_inverse_kinematics.py" +) + + +# --------------------------------------------------------------------------- # +# Helpers to read definitions out of the source without importing it. +# --------------------------------------------------------------------------- # +def _parse_module(path: Path) -> ast.Module: + return ast.parse(path.read_text(), filename=str(path)) + + +def _eval_assigned_dict(module: ast.Module, name: str) -> dict: + """Return the literally-assigned dict for a top-level ``name = {...}``. + + Evaluates value expressions in a tiny sandbox where ``np.deg2rad`` is the real + numpy function (so bounds tuples come out in radians, exactly as the module + would compute them) and ``np`` is numpy. Only used on trusted in-repo source. + """ + for node in module.body: + if isinstance(node, ast.Assign): + targets = [t.id for t in node.targets if isinstance(t, ast.Name)] + if name in targets: + return eval( # noqa: S307 - trusted, in-repo source + compile(ast.Expression(node.value), "", "eval"), + {"np": np, "__builtins__": {}}, + {}, + ) + raise AssertionError(f"Top-level dict assignment {name!r} not found in source") + + +def _extract_function_source(path: Path, func_name: str) -> str: + """Return the exact source text of a top-level function definition.""" + module = _parse_module(path) + src = path.read_text() + for node in module.body: + if isinstance(node, ast.FunctionDef) and node.name == func_name: + return ast.get_source_segment(src, node) + raise AssertionError(f"Function {func_name!r} not found in {path}") + + +def _load_pure_functions(path: Path, func_names: list[str]) -> dict: + """Exec the given functions' real source in an isolated numpy-only namespace. + + A stub ``logger`` (loguru-like, swallows calls) is provided so the functions' + logging calls do not require loguru. This lets us test the genuine helper + bodies without importing the flygym-tainted module. + """ + + class _StubLogger: + def __getattr__(self, _name): + return lambda *a, **k: None + + ns = {"np": np, "logger": _StubLogger()} + for fn in func_names: + exec(_extract_function_source(path, fn), ns) # noqa: S102 - trusted source + return ns + + +# --------------------------------------------------------------------------- # +# I3-A: the deleted DOF map must be restored, with the right keys/order/values. +# --------------------------------------------------------------------------- # +# seqikpy emits joint-angle dict keys "Angle_{leg}_{canonical_dof}" in this order +# (see seqikpy.leg_inverse_kinematics.LegInvKinSeq); the IK save path packs DOFs +# in the order of dof_name_lookup_canonical_to_nmf.keys(), so the keys must be +# exactly these canonical names in this order. +EXPECTED_CANONICAL_DOFS_IN_ORDER = [ + "ThC_yaw", + "ThC_pitch", + "ThC_roll", + "CTr_pitch", + "CTr_roll", + "FTi_pitch", + "TiTa_pitch", +] +# The canonical -> NMF DOF name mapping (flygym-v1 NMF joint names). +EXPECTED_CANONICAL_TO_NMF = { + "ThC_yaw": "Coxa_yaw", + "ThC_pitch": "Coxa", + "ThC_roll": "Coxa_roll", + "CTr_pitch": "Femur", + "CTr_roll": "Femur_roll", + "FTi_pitch": "Tibia", + "TiTa_pitch": "Tarsus1", +} + + +def test_i3a_dof_map_restored_with_correct_keys_and_order(): + module = _parse_module(CONSTANTS_PY) + dof_map = _eval_assigned_dict(module, "dof_name_lookup_canonical_to_nmf") + + # Keys must be exactly the 7 canonical DOFs, in the documented order. + assert list(dof_map.keys()) == EXPECTED_CANONICAL_DOFS_IN_ORDER + # Values must be the matching NMF DOF names. + assert dof_map == EXPECTED_CANONICAL_TO_NMF + + +def test_i3a_dof_keys_match_seqikpy_emitted_angle_keys(): + """The call sites build `Angle_{leg}_{dof}` from these keys; every key must be + a canonical DOF name that seqikpy actually emits (ThC_yaw, ..., TiTa_pitch).""" + module = _parse_module(CONSTANTS_PY) + dof_map = _eval_assigned_dict(module, "dof_name_lookup_canonical_to_nmf") + assert set(dof_map.keys()) == set(EXPECTED_CANONICAL_DOFS_IN_ORDER) + # 7 leg DOFs total (matches the (n_frames, 6, 7) joint-angle array). + assert len(dof_map) == 7 + + +def test_i3a_call_sites_consume_the_restored_constant(): + """Guard against the call sites being changed out from under the constant.""" + for path in (RUN_IK_PY, REPO_ROOT + / "src" / "poseforge" / "production" / "spotlight" / "keypoints3d.py"): + text = path.read_text() + assert "dof_name_lookup_canonical_to_nmf.keys()" in text, ( + f"{path} no longer consumes dof_name_lookup_canonical_to_nmf; " + f"update this test and the constant together." + ) + + +# --------------------------------------------------------------------------- # +# I3-B: nmf_bounds must be L/R mirror-consistent for every leg/DOF. +# --------------------------------------------------------------------------- # +LEG_DOFS = [ + "ThC_yaw", + "ThC_pitch", + "ThC_roll", + "CTr_pitch", + "CTr_roll", + "FTi_pitch", + "TiTa_pitch", +] + + +def _mirror_of_left(dof: str, lo: float, hi: float) -> tuple[float, float]: + """Expected RIGHT bound given the LEFT bound (lo, hi). + + Convention (see constants.py I3-B comment): roll/yaw mirror by negating and + swapping the limits -> R = (-hi, -lo); pitch is shared -> R = (lo, hi). + """ + if dof.endswith("pitch"): + return (lo, hi) + return (-hi, -lo) + + +@pytest.fixture(scope="module") +def nmf_bounds() -> dict: + module = _parse_module(CONSTANTS_PY) + return _eval_assigned_dict(module, "nmf_bounds") + + +def test_i3b_all_legs_dofs_present(nmf_bounds): + for side in "LR": + for pos in "FMH": + for dof in LEG_DOFS: + assert f"{side}{pos}_{dof}" in nmf_bounds + + +@pytest.mark.parametrize("pos", ["F", "M", "H"]) +@pytest.mark.parametrize("dof", LEG_DOFS) +def test_i3b_left_right_bounds_are_mirror_consistent(nmf_bounds, pos, dof): + left = tuple(nmf_bounds[f"L{pos}_{dof}"]) + right = tuple(nmf_bounds[f"R{pos}_{dof}"]) + expected_right = _mirror_of_left(dof, *left) + assert right == pytest.approx(expected_right), ( + f"{pos}_{dof}: right bound {np.rad2deg(right)} deg is not the mirror of " + f"left bound {np.rad2deg(left)} deg (expected " + f"{np.rad2deg(expected_right)} deg)." + ) + + +def test_i3b_no_physically_implausible_lower_bound(nmf_bounds): + """The old RF/RM CTr_pitch lower bound was -270 deg (implausible). After the + fix, no bound should exceed +/-180 deg.""" + for key, (lo, hi) in nmf_bounds.items(): + assert np.rad2deg(lo) >= -180.0 - 1e-6, f"{key} lower bound < -180 deg" + assert np.rad2deg(hi) <= 180.0 + 1e-6, f"{key} upper bound > 180 deg" + + +def test_i3b_specific_fixed_values(nmf_bounds): + """Pin the exact post-fix values for the five corrected bounds.""" + expected_deg = { + "RF_ThC_roll": (-90, 10), + "RF_CTr_pitch": (-180, 10), + "RM_ThC_yaw": (-90, 45), + "RM_CTr_pitch": (-180, 10), + "RH_ThC_yaw": (-90, 45), + } + for key, (lo_deg, hi_deg) in expected_deg.items(): + lo, hi = nmf_bounds[key] + assert (np.rad2deg(lo), np.rad2deg(hi)) == pytest.approx((lo_deg, hi_deg)) + + +# --------------------------------------------------------------------------- # +# I3-D: NaN-frame interpolation helper (real source, numpy-only sandbox). +# --------------------------------------------------------------------------- # +@pytest.fixture(scope="module") +def interp_fn(): + ns = _load_pure_functions(INVKIN_PY, ["_interpolate_nan_frames"]) + return ns["_interpolate_nan_frames"] + + +def test_i3d_no_nan_is_passthrough(interp_fn): + data = np.arange(2 * 5 * 3, dtype=np.float32).reshape(2, 5, 3) + filled, n_nan, n_frames = interp_fn(data) + assert n_nan == 0 + assert n_frames == 0 + assert np.array_equal(filled, data) + + +def test_i3d_interior_gap_is_linearly_interpolated(interp_fn): + # One keypoint, one coord effectively: build (n_frames, 1, 1)-ish via 3 frames. + data = np.zeros((3, 1, 3), dtype=np.float32) + data[:, 0, 0] = [0.0, np.nan, 4.0] # interior NaN -> should become 2.0 + data[:, 0, 1] = [1.0, 1.0, 1.0] + data[:, 0, 2] = [0.0, 2.0, 4.0] + filled, n_nan, n_affected = interp_fn(data) + assert n_nan == 1 + assert n_affected == 1 + assert filled[1, 0, 0] == pytest.approx(2.0) + assert not np.isnan(filled).any() + + +def test_i3d_edge_nans_are_filled_with_nearest(interp_fn): + data = np.zeros((4, 1, 3), dtype=np.float32) + # leading and trailing NaN -> forward/backward fill at edges + data[:, 0, 0] = [np.nan, 2.0, 4.0, np.nan] + filled, n_nan, n_affected = interp_fn(data) + assert n_nan == 2 + assert n_affected == 2 + assert filled[0, 0, 0] == pytest.approx(2.0) # backfilled from first valid + assert filled[3, 0, 0] == pytest.approx(4.0) # forward-filled from last valid + assert not np.isnan(filled).any() + + +def test_i3d_fully_nan_series_stays_nan(interp_fn): + data = np.zeros((3, 1, 3), dtype=np.float32) + data[:, 0, 0] = [np.nan, np.nan, np.nan] # whole series NaN -> unrecoverable + data[:, 0, 1] = [1.0, 2.0, 3.0] + filled, n_nan, n_affected = interp_fn(data) + assert n_nan == 3 + assert n_affected == 3 + # The fully-NaN coordinate cannot be recovered and must remain NaN so the + # caller can detect and raise. + assert np.isnan(filled[:, 0, 0]).all() + # The healthy coordinate is untouched. + assert np.array_equal(filled[:, 0, 1], np.array([1.0, 2.0, 3.0], dtype=np.float32)) + + +def test_i3d_affected_frame_count(interp_fn): + data = np.zeros((4, 2, 3), dtype=np.float32) + data[1, 0, 0] = np.nan + data[1, 1, 2] = np.nan # same frame, two NaNs + data[3, 0, 1] = np.nan # another frame + _, n_nan, n_affected = interp_fn(data) + assert n_nan == 3 + assert n_affected == 2 # frames 1 and 3 + + +# --------------------------------------------------------------------------- # +# I3-C: large frame-to-frame joint-angle jump detector (real source). +# --------------------------------------------------------------------------- # +@pytest.fixture(scope="module") +def jump_fns(): + return _load_pure_functions( + INVKIN_PY, + ["detect_large_joint_angle_jumps", "log_large_joint_angle_jumps"], + ) + + +def test_i3c_detects_jump_above_threshold(jump_fns): + detect = jump_fns["detect_large_joint_angle_jumps"] + series = np.array([0.0, 0.1, 0.2, 2.0, 2.1]) # big jump 0.2 -> 2.0 + masks = detect({"Angle_LF_ThC_yaw": series}, threshold_rad=1.0) + mask = masks["Angle_LF_ThC_yaw"] + assert mask.dtype == bool + assert mask[0] == False # frame 0 never flagged + assert mask[3] == True # the 1.8 rad jump + assert mask.sum() == 1 + + +def test_i3c_no_false_positive_for_smooth_series(jump_fns): + detect = jump_fns["detect_large_joint_angle_jumps"] + series = np.linspace(0.0, 1.0, 50) # max step ~0.02 rad + masks = detect({"Angle_RF_FTi_pitch": series}, threshold_rad=np.deg2rad(45)) + assert masks["Angle_RF_FTi_pitch"].sum() == 0 + + +def test_i3c_skips_non_1d_and_short_series(jump_fns): + detect = jump_fns["detect_large_joint_angle_jumps"] + masks = detect( + { + "scalar": np.array([1.0]), # too short + "twod": np.zeros((3, 2)), # not 1D + "ok": np.array([0.0, 10.0]), # valid, one jump + }, + threshold_rad=1.0, + ) + assert "scalar" not in masks + assert "twod" not in masks + assert masks["ok"].tolist() == [False, True] + + +def test_i3c_log_helper_counts_total_jumps(jump_fns): + log = jump_fns["log_large_joint_angle_jumps"] + total = log( + { + "Angle_LF_ThC_yaw": np.array([0.0, 0.0, 5.0]), # 1 jump + "Angle_LF_ThC_pitch": np.array([0.0, 5.0, 10.0]), # 2 jumps + }, + threshold_rad=1.0, + ) + assert total == 3 + + +def test_i3c_log_helper_zero_when_smooth(jump_fns): + log = jump_fns["log_large_joint_angle_jumps"] + total = log({"a": np.linspace(0, 1, 100)}, threshold_rad=np.deg2rad(45)) + assert total == 0