aind-behavior-vr-foraging-packaging

Aind.Behavior.VrForaging.Packaging

Parses raw AIND VR-foraging behavioral sessions into analysis-ready parquet tables and an NWB file.

Architecture

A session is loaded once (via contraqctor), then a set of independent processors fan out over it. Each processor owns one output and knows how to express it in two targets:

raw session dir
      │
      ▼
  Dataset  ◄── aind_behavior_vr_foraging.data_contract.dataset(path)
      │
      ▼
  create_processors(dataset)          # picks processor variants by dataset version
      │   [SessionMetadata, PositionAndVelocity, SiteTable, Licks, Sniffing,
      │    SoftwareEvents, Events]
      │
      ├─► proc.compute()  ──► pandas DataFrame  ──► one <name>.parquet   (process_session)
      │                        (provenance stamped into df.attrs / parquet schema)
      │
      └─► proc.nwbize(nwb) ──► populates an NWBFile ──► .nwb.zarr (NwbSession)

• Processor — every processor subclasses AbstractProcessor, implementing _compute() and (optionally) nwbize(). compute() wraps _compute() and stamps provenance (packaging_version, data_contract_version, dataset_version, processor) into the DataFrame's attrs.
• DataFrame — the common in-memory representation. One row per unit of the output (e.g. one site-table row = one site).
• Parquet — pipeline.session.process_session() calls compute() on each processor and writes a parquet per processor, promoting df.attrs to first-class parquet metadata (readable from DuckDB, Polars, R arrow, Spark, …).
• NWB — NwbSession builds a single NWBFile from AIND metadata, then calls each processor's nwbize() to fill it, and writes NWB-Zarr.

Version dispatch is automatic: datasets with schema version < 0.6.0 receive legacy processor variants.

Examples

• Walkthrough of the parquet workflows (all-at-once, single stream, load-back): docs/guides/session-from-disk.md
• Query the local export with pandas and DuckDB: docs/examples/query_export.py
• Query from S3 with DuckDB: docs/examples/query_export_s3.py
• Query from S3 with Polars: docs/examples/query_export_s3_polars.py
• Full architecture docs: docs/knowledge/ (start at overview.md)

Get a sites table

Install straight from GitHub with uv:

# into a uv project
uv add "git+https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging.git"

# or into the current environment
uv pip install "git+https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging.git"

Then load a session and compute the sites table (one row per site):

from aind_behavior_vr_foraging.data_contract import dataset
from aind_behavior_vr_foraging_packaging.pipeline.session import resolve_site_table_processor

ds = dataset("path/to/session")  # load the raw session
sites_df = resolve_site_table_processor(ds).compute()

sites_df.to_parquet("sites.parquet")  # optional: persist to disk
print(f"{len(sites_df)} sites, {sites_df['has_reward'].sum()} rewarded")

resolve_site_table_processor automatically picks the current or legacy variant based on the dataset's schema version. To produce every table at once, use process_session(ds, "output_dir") instead — it writes sites.parquet, position_velocity.parquet, and the rest, and returns them keyed by name.

Exporting a dataset collection

Install the CLI with uv:

uv tool install aind-behavior-vr-foraging-packaging

Then run the export pipeline across a folder of raw session directories (--input-dir must contain one subdirectory per session):

vr-foraging-packaging batch --input-dir /data/raw --output-dir /data/export

--output-dir receives the results:

/data/export/
├── session.parquet          # session catalogue (one row per session)
├── sites.parquet            # aggregated sites table (all sessions)
└── sessions/
    └── <session_id>/
        ├── sites.parquet
        ├── position_velocity.parquet
        └── ...

Subcommands

Swap batch for session or aggregate in the command above to run those.

The aggregated session.parquet and sites.parquet are sorted by session_id (within-session row order is kept) and written in bounded row groups with min/max statistics. Readers that filter by session, such as pl.scan_parquet("sites.parquet").filter(pl.col("session_id").is_in(ids)), fetch only the row groups that hold those sessions, not the whole file. Run vr-foraging-packaging <command> --help for that command's full flag reference, or see the Getting Started guide for common flag combinations.

Want to try the latest main without installing? uvx runs it straight from GitHub — since the command name differs from the package name, pass the source via --from:

uvx --from "git+https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging.git" \
    vr-foraging-packaging batch --input-dir /data/raw --output-dir /data/export

Documentation

The full documentation site is built with Zensical.

Preview locally:

uv sync --group docs
uv run zensical serve

Build a static copy:

uv run zensical build --clean
# output → site/

The site deploys automatically to GitHub Pages on every push to main as part of the main CI workflow.

Contributors

Contributions to this repository are welcome! However, please ensure that your code adheres to the recommended DevOps practices below:

Linting

We use ruff as our primary linting tool.

Type checking

We use ty for static type checking. Run uv run ty check from the root of the repository.

Testing

Attempt to add tests when new features are added. To run the currently available tests, run uv run pytest from the root of the repository.

Integration tests

Integration tests run the parser end-to-end against real datasets stored in a public S3 bucket. They are gated by a pytest marker so they don't run by default.

Run locally:

uv run pytest -m integration

The first run downloads datasets (~100 MB per dataset) to tests/integration/.cache/. Subsequent runs reuse the cache when the S3 ETag matches. The cache directory is gitignored.

[!IMPORTANT] On Windows, enable long paths first. test_full_pipeline writes an NWB-Zarr file whose chunk paths exceed the legacy 260-character MAX_PATH limit, and it fails with FileNotFoundError: ... .zarray.<hash>.partial — which looks like a parsing bug but is not. Enable long paths once, in an elevated PowerShell, then restart your shell:

New-ItemProperty -Path "HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem" `
  -Name LongPathsEnabled -Value 1 -PropertyType DWORD -Force

If you cannot elevate, uv run pytest -m integration --basetemp=C:\t works around it by shortening the temp path. Linux and macOS are unaffected, as is CI (the integration job runs on ubuntu-latest).

Trigger on a PR:

Integration tests do not run on every PR. To run them for a specific PR, add the run-integration label via the GitHub UI (open the PR, click Labels in the right-hand sidebar, and select run-integration) or with:

gh pr edit <PR_NUMBER> --add-label run-integration

The integration job runs automatically on push to main and on release: published. A release cannot ship without the integration suite passing.

Adding a dataset:

Add an entry to tests/integration/datasets.yml. The manifest schema and full field documentation are in tests/integration/model.py (Pydantic model). The rationale field is required and is printed alongside any test failure to make triage fast.

Lock files

We use uv to manage our lock files and therefore encourage everyone to use uv as a package manager as well.