Add NumPy window-processing engine

[gchindemi]

Summary

• Add selectable xarray and NumPy engines to window selectors and datasets.
• Preserve xarray as the public default.
• Route training, evaluation, and inference through NumPy.
• Support both containers transparently in pose augmentations and PoseToTensor.
• Add selector, dataset, augmentation, and pipeline parity tests.
• Adjust the worker heuristic from 8 to 16 samples per worker.

Performance

On CalMS21 task 1 with all four pretraining tasks:

Engine	Warm epoch	CPU time	GPU utilization
xarray	45.70 s	682.1 s	77.6%
NumPy	35.95 s	47.3 s	99.3%

This reduced warm epoch time by 21.3% and CPU work by approximately 14.4x.

Verification

• Ruff passes.
• Documentation builds successfully.
• Current suite: 305 passed, with the same 24 pre-existing failures observed on the clean baseline.
• Added 154 passing tests without introducing regressions.

Closes #93.

Follow-up: #92 tracks the existing TemporalOrder sampling ambiguity.

[Copilot]

Pull request overview

This PR introduces a selectable NumPy-based window-processing engine alongside the existing xarray path, keeping xarray as the public default while routing internal training/evaluation/inference pipelines through NumPy for substantially lower CPU overhead.

Changes:

• Add engine={"xarray","numpy"} to public selectors/datasets, with NumPy producing canonical (time, individuals, keypoints, space) arrays.
• Update augmentation + preprocessing transforms to transparently accept either xarray datasets or canonical NumPy arrays (with PoseToVideo remaining xarray-only).
• Add extensive parity tests and update documentation; adjust the training DataLoader worker heuristic (via estimate_num_workers) to target 16 samples/worker.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_training_helpers.py	Adds coverage for estimate_num_workers and adapts dataloader tests for the new worker-selection path.
tests/test_numpy_engine.py	Comprehensive parity tests across engines for selectors, datasets, transforms, and internal pipeline routing.
src/lisbet/transforms_extra.py	Adds canonical NumPy handling and dispatch in pose transforms, plus xarray-only enforcement for PoseToVideo.
src/lisbet/training/utils.py	Updates worker heuristic default (batch_size_per_worker=16) and exposes helper used by training.
src/lisbet/training/tasks.py	Routes configured training/dev datasets through the NumPy engine internally.
src/lisbet/training/core.py	Switches dataloader worker selection to the updated estimate_num_workers default behavior.
src/lisbet/inference/common.py	Forces inference WindowDataset to use engine="numpy".
src/lisbet/evaluation.py	Forces evaluation AnnotatedWindowDataset to use engine="numpy".
src/lisbet/datasets/map_style.py	Adds engine argument + tracking to map-style datasets and passes it through selectors.
src/lisbet/datasets/iterable_style.py	Adds engine argument and implements NumPy-compatible concatenation paths for iterable datasets.
src/lisbet/datasets/common.py	Implements NumPy window selection (padding + interpolation) and annotation caching for NumPy engine.
src/lisbet/datasets/init.py	Exports WindowSelector / AnnotatedWindowSelector publicly.
src/lisbet/config/schemas.py	Updates augmentation config documentation to reflect container/semantics changes.
src/lisbet/cli/commands/train.py	Updates CLI help text to reflect updated augmentation semantics.
docs/user_guide/data_preparation.rst	Documents the new window-processing engines and canonical NumPy representation.
docs/user_guide/data_augmentation.rst	Documents container support (xarray vs NumPy) and updates augmentation descriptions/examples.

Comments suppressed due to low confidence (1)

src/lisbet/inference/common.py:197

• The PR description mentions adjusting the worker heuristic from 8→16 samples per worker, but inference still uses batch_size // 8, so large inference batches can still over-provision DataLoader workers. Consider matching the updated heuristic here (or centralizing on a shared helper) to keep behavior consistent with the change rationale.

        engine="numpy",
    )

    num_workers = min(suggested_max_num_workers(1), batch_size // 8)
    prefetch_factor = 4 if num_workers > 0 else None

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.