experiments/analysis/modality_viz.py

"""Visualize mocap skeleton frames, IMU waveforms, EMG waveforms."""
import os, numpy as np, pandas as pd, matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa

REC = "${PULSE_ROOT}/dataset/v1/s1"
OUT = "${PULSE_ROOT}/paper/figures"
os.makedirs(OUT, exist_ok=True)

# ---- Skeleton bone definition (marker pairs) ----
BONES = [
    # torso
    ("HeadTop","HeadFront"),("HeadL","HeadR"),("HeadFront","SpineTop"),
    ("SpineTop","Chest"),("Chest","WaistLFront"),("Chest","WaistRFront"),
    ("WaistLFront","WaistLBack"),("WaistRFront","WaistRBack"),
    ("WaistLBack","BackL"),("WaistRBack","BackR"),("BackL","BackR"),
    ("SpineTop","LShoulderTop"),("SpineTop","RShoulderTop"),
    ("LShoulderTop","LShoulderBack"),("RShoulderTop","RShoulderBack"),
    # left arm
    ("LShoulderTop","LArm"),("LArm","LElbowOut"),("LElbowOut","LElbowBack"),
    ("LElbowOut","LForearmRoll"),("LForearmRoll","LWristOut"),
    ("LWristOut","LWristIn"),("LWristOut","LHandOut"),("LWristIn","LHandIn"),
    ("LHandOut","LIndex2"),("LIndex2","LIndexTip"),
    ("LHandOut","LMiddle2"),("LMiddle2","LMiddleTip"),
    ("LHandIn","LRing2"),("LRing2","LRingTip"),
    ("LHandIn","LPinky2"),("LPinky2","LPinkyTip"),
    ("LWristIn","LThumb1"),("LThumb1","LThumbTip"),
    # right arm
    ("RShoulderTop","RArm"),("RArm","RElbowOut"),("RElbowOut","RElbowBack"),
    ("RElbowOut","RForearmRoll"),("RForearmRoll","RWristOut"),
    ("RWristOut","RWristIn"),("RWristOut","RHandOut"),("RWristIn","RHandIn"),
    ("RHandOut","RIndex2"),("RIndex2","RIndexTip"),
    ("RHandOut","RMiddle2"),("RMiddle2","RMiddleTip"),
    ("RHandIn","RRing2"),("RRing2","RRingTip"),
    ("RHandIn","RPinky2"),("RPinky2","RPinkyTip"),
    ("RWristIn","RThumb1"),("RThumb1","RThumbTip"),
]


def load_mocap(path):
    df = pd.read_csv(path)
    # Extract x,y,z for each marker ignoring Type cols
    markers = {}
    for col in df.columns:
        if col.startswith("Q_") and col.endswith(" X"):
            name = col[2:-2]
            xs = df[f"Q_{name} X"].to_numpy()
            ys = df[f"Q_{name} Y"].to_numpy()
            zs = df[f"Q_{name} Z"].to_numpy()
            markers[name] = np.stack([xs, ys, zs], axis=-1)
    return df["Time"].to_numpy(), markers


def plot_skeletons():
    t, mk = load_mocap(os.path.join(REC, "aligned_mocap_100hz.csv"))
    N = len(t)
    # pick 4 time frames well spread through the recording with valid data
    candidate = np.linspace(int(0.1*N), int(0.9*N), 4).astype(int)

    fig = plt.figure(figsize=(12, 3.2))
    for i, fr in enumerate(candidate):
        ax = fig.add_subplot(1, 4, i+1, projection='3d')
        # gather all points at this frame
        pts = np.array([mk[n][fr] for n in mk])
        pts = pts[~np.isnan(pts).any(axis=1)]
        if len(pts) == 0:
            continue
        # draw bones
        for a, b in BONES:
            if a in mk and b in mk:
                pa, pb = mk[a][fr], mk[b][fr]
                if np.isnan(pa).any() or np.isnan(pb).any():
                    continue
                ax.plot([pa[0], pb[0]], [pa[1], pb[1]], [pa[2], pb[2]],
                        color='#2266aa', lw=1.2)
        ax.scatter(pts[:, 0], pts[:, 1], pts[:, 2], s=4, c='#cc3333', alpha=0.8)
        # equal aspect
        c = pts.mean(0)
        r = np.ptp(pts, axis=0).max() / 2
        ax.set_xlim(c[0]-r, c[0]+r); ax.set_ylim(c[1]-r, c[1]+r); ax.set_zlim(c[2]-r, c[2]+r)
        ax.set_xticks([]); ax.set_yticks([]); ax.set_zticks([])
        ax.set_title(f"t={t[fr]:.1f}s", fontsize=9)
        ax.view_init(elev=12, azim=-75)
    fig.suptitle("MoCap skeleton frames (56-marker Qualisys, v1/s1)", fontsize=11)
    fig.tight_layout()
    out = os.path.join(OUT, "mocap_skeleton.pdf")
    fig.savefig(out, bbox_inches='tight'); fig.savefig(out.replace('.pdf', '.png'), dpi=150, bbox_inches='tight')
    plt.close(fig)
    print("Saved", out)


def plot_imu():
    df = pd.read_csv(os.path.join(REC, "aligned_imu_100hz.csv"))
    t = df["time"].to_numpy(); t = t - t[0]
    # pick 5 body locations (WT0..WT9 order roughly: wrists, forearms, upper arms, shins, thighs, torso)
    sites = [("WT0", "Wrist R"), ("WT2", "Forearm R"),
             ("WT4", "Upper arm R"), ("WT6", "Shin R"), ("WT9", "Torso")]
    fig, axes = plt.subplots(len(sites), 1, figsize=(9, 6), sharex=True)
    # crop to 20s window mid-recording
    mid = len(t)//2
    sl = slice(max(0, mid-1000), min(len(t), mid+1000))
    for ax, (sid, lbl) in zip(axes, sites):
        for comp, col in zip(["x", "y", "z"], ["#d62728", "#2ca02c", "#1f77b4"]):
            ax.plot(t[sl], df[f"{sid}_acc_{comp}"].to_numpy()[sl], color=col, lw=0.8, label=f"acc_{comp}")
        ax.set_ylabel(lbl, fontsize=9)
        ax.grid(alpha=0.3)
    axes[0].legend(loc="upper right", ncol=3, fontsize=8)
    axes[-1].set_xlabel("Time (s)")
    fig.suptitle("IMU 3-axis acceleration across 5 body sites (v1/s1, 20s window)", fontsize=11)
    fig.tight_layout()
    out = os.path.join(OUT, "imu_waveforms.pdf")
    fig.savefig(out, bbox_inches='tight'); fig.savefig(out.replace('.pdf', '.png'), dpi=150, bbox_inches='tight')
    plt.close(fig)
    print("Saved", out)


def plot_emg():
    df = pd.read_csv(os.path.join(REC, "aligned_emg_100hz.csv"))
    t = df["time"].to_numpy(); t = t - t[0]
    ch = [f"emg_{i}" for i in range(1, 9)]
    # 20s window mid-recording
    mid = len(t)//2
    sl = slice(max(0, mid-1000), min(len(t), mid+1000))
    fig, axes = plt.subplots(8, 1, figsize=(9, 7), sharex=True)
    for ax, c in zip(axes, ch):
        sig = df[c].to_numpy()[sl]
        ax.plot(t[sl], sig, color="#555", lw=0.5)
        # envelope overlay
        env = pd.Series(np.abs(sig)).rolling(20, min_periods=1).mean().to_numpy()
        ax.plot(t[sl], env, color="#d62728", lw=0.9)
        ax.set_ylabel(c, fontsize=8)
        ax.grid(alpha=0.3)
    axes[-1].set_xlabel("Time (s)")
    fig.suptitle("Surface EMG 8-channel raw (grey) with rectified envelope (red), v1/s1, 20s window",
                 fontsize=11)
    fig.tight_layout()
    out = os.path.join(OUT, "emg_waveforms.pdf")
    fig.savefig(out, bbox_inches='tight'); fig.savefig(out.replace('.pdf', '.png'), dpi=150, bbox_inches='tight')
    plt.close(fig)
    print("Saved", out)


if __name__ == "__main__":
    plot_skeletons()
    plot_imu()
    plot_emg()