example_load_dataset.py

#!/usr/bin/env python3
"""
Minimal workable example for the ClothTransformer cloth-simulation dataset.

What it does
------------
1. Loads one simulation `.npz` sample.
2. Prints a summary of every array (shape / dtype / meaning).
3. Shows how to convert the stored per-frame displacements into physical
   velocities (world-units per second).
4. Reconstructs the per-triangle collider geometry from the raw arrays.
5. Exports a single frame of the cloth mesh and the collider mesh to OBJ
   files so you can open them in any 3D viewer.

Requirements: numpy only.

Usage
-----
    python example_load_dataset.py
    python example_load_dataset.py --npz sims_grasp_all_case1/sim_00007.npz --frame 120
"""

import argparse
import os

import numpy as np

# Simulation constants (see README.md).
DT = 1.0 / 60.0          # seconds per frame
N_FRAMES = 240           # frames per trajectory


def load_sample(path):
    """Load one .npz sample into a plain dict of numpy arrays."""
    with np.load(path) as data:
        return {k: data[k] for k in data.files}


def summarize(sample):
    """Print shapes / dtypes and derive the basic dimensions."""
    print("Arrays in this sample:")
    for k, v in sample.items():
        print(f"  {k:22s} shape={str(v.shape):20s} dtype={v.dtype}")

    n_v = sample["traj"].shape[1]                # cloth vertices
    n_c = sample["collision_vertices"].shape[1]  # collider vertices
    n_frames = sample["traj"].shape[0]
    print(f"\nCloth   : {n_v} vertices, {len(sample['triangles'])} triangles")
    print(f"Collider: {n_c} vertices, {len(sample['collision_triangles'])} triangles")
    print(f"Frames  : {n_frames} (dt = {DT:.6f} s, duration = {n_frames * DT:.3f} s)")


def cloth_rest_pose_and_uv(sample):
    """Split the `initial` field into rest positions and UV coordinates."""
    initial = sample["initial"]              # (N_v, 6), float64
    rest_pos = initial[:, 0:3]               # 3D rest position (== traj[0])
    uv = initial[:, 3:5]                     # (u, v); the 6th column is always 0
    assert np.allclose(rest_pos, sample["traj"][0]), "rest_pos should equal frame 0"
    return rest_pos, uv


def physical_velocity(displacement):
    """Stored arrays are per-frame displacement; divide by DT for units/second."""
    return displacement / DT


def reconstruct_collider_triangles(sample):
    """Gather collider vertex positions into per-triangle (T, F_col, 9) form.

    This reproduces the convenience `collision` field from the raw arrays.
    """
    cv = sample["collision_vertices"]        # (T, N_c, 3)
    tri = sample["collision_triangles"]      # (F_col, 3)
    return cv[:, tri.reshape(-1)].reshape(cv.shape[0], -1, 9)


def write_obj(path, verts, faces):
    """Write a triangle mesh (verts: (N,3), faces: (F,3), 0-based) to an OBJ file."""
    with open(path, "w") as f:
        for x, y, z in verts:
            f.write(f"v {x:.6f} {y:.6f} {z:.6f}\n")
        for a, b, c in faces:                # OBJ is 1-based
            f.write(f"f {a + 1} {b + 1} {c + 1}\n")
    print(f"  wrote {path}  ({len(verts)} verts, {len(faces)} faces)")


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--npz",
        default=os.path.join(here, "sims_collision_all_static_1k", "sim_00000.npz"),
        help="Path to a sim_*.npz file (absolute, or relative to this folder).",
    )
    parser.add_argument("--frame", type=int, default=120, help="Frame index to export.")
    parser.add_argument("--out", default=here, help="Directory for exported OBJ files.")
    args = parser.parse_args()

    npz_path = args.npz if os.path.isabs(args.npz) else os.path.join(here, args.npz)
    print(f"Loading: {npz_path}\n")
    sample = load_sample(npz_path)

    # 1. Summary --------------------------------------------------------------
    summarize(sample)

    # 2. Rest pose + UVs ------------------------------------------------------
    rest_pos, uv = cloth_rest_pose_and_uv(sample)
    print(f"\nUV bounds: u in [{uv[:, 0].min():.3f}, {uv[:, 0].max():.3f}], "
          f"v in [{uv[:, 1].min():.3f}, {uv[:, 1].max():.3f}]")

    # 3. Displacement -> physical velocity ------------------------------------
    f = max(1, min(args.frame, sample["traj"].shape[0] - 1))
    cloth_vel = physical_velocity(sample["traj_vel"][f])         # (N_v, 3) units/s
    print(f"\nFrame {f}: mean cloth speed = "
          f"{np.linalg.norm(cloth_vel, axis=1).mean():.4f} units/s")

    # Sanity check: stored displacement equals the position difference.
    assert np.allclose(
        sample["traj_vel"][f], sample["traj"][f] - sample["traj"][f - 1]
    ), "traj_vel[t] should equal traj[t] - traj[t-1]"

    # 4. Reconstruct collider triangles (matches the stored `collision` field)
    recon = reconstruct_collider_triangles(sample)
    if "collision" in sample:
        ok = np.allclose(recon, sample["collision"])
        print(f"Reconstructed collider triangles match `collision` field: {ok}")

    # 5. Export one frame to OBJ ----------------------------------------------
    print(f"\nExporting frame {f} to OBJ:")
    write_obj(os.path.join(args.out, f"cloth_frame{f:03d}.obj"),
              sample["traj"][f], sample["triangles"])
    write_obj(os.path.join(args.out, f"collider_frame{f:03d}.obj"),
              sample["collision_vertices"][f], sample["collision_triangles"])

    print("\nDone.")


if __name__ == "__main__":
    main()