aind-behavior-vr-foraging-packaging 0.0.3__tar.gz

aind_behavior_vr_foraging_packaging-0.0.3/PKG-INFO

@@ -0,0 +1,156 @@
+Metadata-Version: 2.4
+Name: aind-behavior-vr-foraging-packaging
+Version: 0.0.3
+Summary: Parses and packages AIND VR Foraging behavioral sessions into tabular (parquet) and NWB outputs.
+Author: Bruno Cruz
+Author-email: Bruno Cruz <bruno.cruz@alleninstitute.org>
+License-Expression: MIT
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Operating System :: Microsoft :: Windows
+Requires-Dist: aind-behavior-vr-foraging[data]>=1
+Requires-Dist: aind-data-schema>2.7
+Requires-Dist: hdmf-zarr
+Requires-Dist: numpy>=2
+Requires-Dist: pandas
+Requires-Dist: pydantic-settings>=2.10.1
+Requires-Dist: pydantic>=2.10
+Requires-Dist: pynwb
+Requires-Dist: scipy
+Requires-Dist: pydantic-settings
+Requires-Dist: aind-nwb-utils
+Requires-Dist: ndx-events
+Requires-Dist: aind-data-access-api>=1.9.2
+Requires-Python: >=3.11
+Project-URL: Documentation, https://allenneuraldynamics.github.io/Aind.Behavior.VrForaging.Packaging/
+Project-URL: Repository, https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/
+Project-URL: Issues, https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/issues
+Project-URL: Changelog, https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/releases
+Description-Content-Type: text/markdown
+
+# Aind.Behavior.VrForaging.Packaging
+
+![CI](https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/actions/workflows/aind-behavior-vr-foraging-packaging.yml/badge.svg)
+[![License](https://img.shields.io/badge/license-MIT-brightgreen)](LICENSE)
+[![ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+
+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:
+
+```text
+raw session dir
+      │
+      ▼
+  Dataset  ◄── aind_behavior_vr_foraging.data_contract.dataset(path)
+      │
+      ▼
+  create_processors(dataset)          # picks processor variants by dataset version
+      │   [TrialTable, PositionAndVelocity, Licks, Sniffing, SoftwareEvents, Events]
+      │
+      ├─► proc.compute()  ──► pandas DataFrame  ──► one <name>.parquet   (run_session)
+      │                        (provenance stamped into df.attrs / parquet schema)
+      │
+      └─► proc.nwbize(nwb) ──► populates an NdxEventsNWBFile ──► .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 trial-table row = one *site*).
+- **Parquet** — `pipeline.run_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 `NdxEventsNWBFile` 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
+
+- Runnable script covering the parquet workflows (all-at-once, single stream,
+  load-back): [scripts/example_parquet_pipeline.py](scripts/example_parquet_pipeline.py)
+- The NWB workflow: [docs/knowledge/architecture/nwb-packaging.md](docs/knowledge/architecture/nwb-packaging.md)
+- Full architecture docs: [docs/knowledge/](docs/knowledge/) (start at [overview.md](docs/knowledge/overview.md))
+
+### Get a trials table
+
+Install straight from GitHub with [uv](https://docs.astral.sh/uv/):
+
+```bash
+# 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 trials table (one row per *site*):
+
+```python
+from aind_behavior_vr_foraging.data_contract import dataset
+from aind_behavior_vr_foraging_packaging.pipeline import get_trial_table_processor
+
+ds = dataset("path/to/session")           # load the raw session
+trials_df = get_trial_table_processor(ds).compute()
+
+trials_df.to_parquet("trials.parquet")    # optional: persist to disk
+print(f"{len(trials_df)} sites, {trials_df['has_reward'].sum()} rewarded")
+```
+
+`get_trial_table_processor` automatically picks the current or legacy variant
+based on the dataset's schema version. To produce every table at once, use
+`run_session(ds, "output_dir")` instead — it writes `trials.parquet`,
+`position_velocity.parquet`, and the rest, and returns them keyed by name.
+
+## Contributors
+
+Contributions to this repository are welcome! However, please ensure that your code adheres to the recommended DevOps practices below:
+
+### Linting
+
+We use [ruff](https://docs.astral.sh/ruff/) as our primary linting tool.
+
+### 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:**
+
+```bash
+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.
+
+**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:
+
+```bash
+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](https://docs.astral.sh/uv/) to manage our lock files and therefore encourage everyone to use uv as a package manager as well.

aind_behavior_vr_foraging_packaging-0.0.3/README.md

@@ -0,0 +1,125 @@
+# Aind.Behavior.VrForaging.Packaging
+
+![CI](https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/actions/workflows/aind-behavior-vr-foraging-packaging.yml/badge.svg)
+[![License](https://img.shields.io/badge/license-MIT-brightgreen)](LICENSE)
+[![ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+
+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:
+
+```text
+raw session dir
+      │
+      ▼
+  Dataset  ◄── aind_behavior_vr_foraging.data_contract.dataset(path)
+      │
+      ▼
+  create_processors(dataset)          # picks processor variants by dataset version
+      │   [TrialTable, PositionAndVelocity, Licks, Sniffing, SoftwareEvents, Events]
+      │
+      ├─► proc.compute()  ──► pandas DataFrame  ──► one <name>.parquet   (run_session)
+      │                        (provenance stamped into df.attrs / parquet schema)
+      │
+      └─► proc.nwbize(nwb) ──► populates an NdxEventsNWBFile ──► .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 trial-table row = one *site*).
+- **Parquet** — `pipeline.run_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 `NdxEventsNWBFile` 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
+
+- Runnable script covering the parquet workflows (all-at-once, single stream,
+  load-back): [scripts/example_parquet_pipeline.py](scripts/example_parquet_pipeline.py)
+- The NWB workflow: [docs/knowledge/architecture/nwb-packaging.md](docs/knowledge/architecture/nwb-packaging.md)
+- Full architecture docs: [docs/knowledge/](docs/knowledge/) (start at [overview.md](docs/knowledge/overview.md))
+
+### Get a trials table
+
+Install straight from GitHub with [uv](https://docs.astral.sh/uv/):
+
+```bash
+# 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 trials table (one row per *site*):
+
+```python
+from aind_behavior_vr_foraging.data_contract import dataset
+from aind_behavior_vr_foraging_packaging.pipeline import get_trial_table_processor
+
+ds = dataset("path/to/session")           # load the raw session
+trials_df = get_trial_table_processor(ds).compute()
+
+trials_df.to_parquet("trials.parquet")    # optional: persist to disk
+print(f"{len(trials_df)} sites, {trials_df['has_reward'].sum()} rewarded")
+```
+
+`get_trial_table_processor` automatically picks the current or legacy variant
+based on the dataset's schema version. To produce every table at once, use
+`run_session(ds, "output_dir")` instead — it writes `trials.parquet`,
+`position_velocity.parquet`, and the rest, and returns them keyed by name.
+
+## Contributors
+
+Contributions to this repository are welcome! However, please ensure that your code adheres to the recommended DevOps practices below:
+
+### Linting
+
+We use [ruff](https://docs.astral.sh/ruff/) as our primary linting tool.
+
+### 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:**
+
+```bash
+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.
+
+**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:
+
+```bash
+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](https://docs.astral.sh/uv/) to manage our lock files and therefore encourage everyone to use uv as a package manager as well.

aind_behavior_vr_foraging_packaging-0.0.3/pyproject.toml

@@ -0,0 +1,83 @@
+[build-system]
+requires = ["uv_build>=0.8.22"]
+build-backend = "uv_build"
+
+[project]
+name = "aind-behavior-vr-foraging-packaging"
+description = "Parses and packages AIND VR Foraging behavioral sessions into tabular (parquet) and NWB outputs."
+authors = [
+    {name = "Bruno Cruz", email = "bruno.cruz@alleninstitute.org"},
+    ]
+license = "MIT"
+version = "0.0.3"
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Operating System :: Microsoft :: Windows",
+]
+readme = {file = "README.md", content-type = "text/markdown"}
+
+dependencies = [
+    "aind-behavior-vr-foraging[data] >= 1",
+    "aind-data-schema>2.7",
+    "hdmf-zarr",
+    "numpy>=2",
+    "pandas",
+    "pydantic-settings>=2.10.1",
+    "pydantic>=2.10",
+    "pynwb",
+    "scipy",
+    "pydantic-settings",
+    "aind-nwb-utils",
+    "ndx-events",
+    "aind-data-access-api>=1.9.2",
+]
+
+[project.urls]
+Documentation = "https://allenneuraldynamics.github.io/Aind.Behavior.VrForaging.Packaging/"
+Repository = "https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/"
+Issues = "https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/issues"
+Changelog = "https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging.Packaging/releases"
+
+[dependency-groups]
+
+dev = [
+    'ruff',
+    'pytest',
+    'pytest-cov',
+    'codespell',
+    'ty',
+    "pyarrow>=24.0.0",
+]
+
+[tool.uv]
+default-groups = ['dev']
+
+[tool.ruff]
+line-length = 120
+target-version = 'py311'
+
+[tool.ruff.lint]
+extend-select = ['Q', 'RUF100', 'C90', 'I']
+extend-ignore = []
+mccabe = { max-complexity = 14 }
+pydocstyle = { convention = 'google' }
+
+[tool.codespell]
+skip = '.git,*.pdf,*.svg,uv.lock'
+ignore-words-list = 'nd'
+
+[tool.pytest.ini_options]
+addopts = "--strict-markers --tb=short --cov=src --cov-report=term-missing --cov-fail-under=0 -m 'not integration'"
+markers = [
+  "integration: integration tests that download data from S3 (run with `-m integration`)",
+]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+
+[project.scripts]
+curriculum = "aind_behavior_vr_foraging_packaging.cli:main"

aind_behavior_vr_foraging_packaging-0.0.3/src/aind_behavior_vr_foraging_packaging/__init__.py

@@ -0,0 +1,30 @@
+import re
+from importlib.metadata import PackageNotFoundError, version
+
+
+def pep440_to_semver(ver: str) -> str:
+    """
+    Convert a PEP 440 version to a SemVer-compatible string.
+
+    Examples:
+        1.2.3rc2   -> 1.2.3-rc2
+        1.2.3a1    -> 1.2.3-a1
+        1.2.3b1    -> 1.2.3-b1
+        1.2.3.dev4 -> 1.2.3-dev4
+        1.2.3.post1 -> 1.2.3+post1
+    """
+    # pre-release: a, b, rc -> -aN, -bN, -rcN
+    ver = re.sub(r"(?<=\d)(a|b|rc)(\d+)", r"-\1\2", ver)
+    # dev release: .devN -> -devN
+    ver = re.sub(r"\.dev(\d+)", r"-dev\1", ver)
+    # post release: .postN -> +postN
+    ver = re.sub(r"\.post(\d+)", r"+post\1", ver)
+    return ver
+
+
+try:
+    __version__ = version(__name__)
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+__semver__ = pep440_to_semver(__version__)

aind_behavior_vr_foraging_packaging-0.0.3/src/aind_behavior_vr_foraging_packaging/_base.py

@@ -0,0 +1,98 @@
+import abc
+import re
+import typing as ty
+
+import aind_behavior_vr_foraging
+import pandas as pd
+import semver
+from contraqctor.contract import Dataset
+
+
+def _class_name_to_snake(name: str) -> str:
+    """Convert a CamelCase class name to snake_case, e.g. ``LicksProcessor`` → ``licks_processor``."""
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
+
+
+class AbstractProcessor(abc.ABC):
+    #: Override in subclasses to set a canonical parquet filename stem (e.g. ``"trials"``).
+    #: When ``None`` (the default), ``output_name`` falls back to a snake_case of the class name.
+    __output_name__: ty.ClassVar[str | None] = None
+
+    @property
+    def output_name(self) -> str:
+        """Canonical name used as the parquet filename stem.
+
+        Returns ``__output_name__`` if defined on the class, otherwise a
+        snake_case of the class name (e.g. ``LicksProcessor`` → ``licks_processor``).
+        """
+        return self.__class__.__output_name__ or _class_name_to_snake(type(self).__name__)
+
+    def __init__(self, dataset: Dataset, *, raise_on_error: bool = False) -> None:
+        self._dataset = dataset
+        self._raise_on_error = raise_on_error
+
+    @property
+    def dataset(self) -> Dataset:
+        return self._dataset
+
+    @property
+    def dataset_version(self) -> semver.Version:
+        return self._parse_version(self.dataset.version)
+
+    @property
+    def parser_version(self) -> semver.Version:
+        return semver.Version.parse(aind_behavior_vr_foraging.__semver__)
+
+    @staticmethod
+    def _parse_version(value: str | semver.Version) -> semver.Version:
+        if isinstance(value, semver.Version):
+            return value
+        return semver.Version.parse(value)
+
+    @abc.abstractmethod
+    def _compute(self) -> pd.DataFrame:
+        """Compute this processor's output as a DataFrame.
+
+        Subclasses implement this method. Callers should use :meth:`compute`,
+        which wraps ``_compute`` and stamps provenance metadata into ``df.attrs``.
+        """
+        raise NotImplementedError
+
+    def compute(self) -> pd.DataFrame:
+        """Return the processor's output DataFrame with provenance metadata in attrs.
+
+        Calls :meth:`_compute`, then stamps ``df.attrs`` with:
+
+        - ``packaging_version``: version of this package (``aind-behavior-vr-foraging-packaging``)
+        - ``data_contract_version``: version of ``aind-behavior-vr-foraging`` (defines the behavioral data schema)
+        - ``dataset_version``: actual version recorded in the session's ``tasklogic_input.json``
+        - ``processor``: this processor's class name
+
+        Attrs already set by ``_compute`` (e.g. ``sampling_rate_hz`` from
+        :class:`SniffingProcessor`) are preserved via ``setdefault``.
+        """
+        from importlib.metadata import version as _pkg_version
+
+        df = self._compute()
+        df.attrs.setdefault("packaging_version", _pkg_version("aind-behavior-vr-foraging-packaging"))
+        df.attrs.setdefault("data_contract_version", str(self.parser_version))
+        df.attrs.setdefault("dataset_version", str(self.dataset_version))
+        df.attrs.setdefault("processor", type(self).__name__)
+        return df
+
+    def nwbize(self, nwb_file: ty.Any) -> ty.Any:
+        """Write this processor's output to *nwb_file* and return it.
+
+        Default implementation is a no-op. Override in subclasses that have
+        an NWB representation. May call ``compute()`` internally; the two
+        methods are intentionally independent (no shared state).
+        """
+        return nwb_file
+
+    def with_raise_errors(self, raise_on_error: bool = True) -> ty.Self:
+        self._raise_on_error = raise_on_error
+        return self
+
+    @property
+    def raise_on_error(self) -> bool:
+        return self._raise_on_error

aind_behavior_vr_foraging_packaging-0.0.3/src/aind_behavior_vr_foraging_packaging/acquisition/__init__.py

@@ -0,0 +1,94 @@
+import logging
+import typing as ty
+
+import contraqctor.contract as data_contract
+import pandas as pd
+
+from .._base import AbstractProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class AcquisitionProcessor(AbstractProcessor):
+    __output_name__ = "acquisition"
+
+    def __init__(self, dataset: data_contract.Dataset, *, raise_on_error: bool = False) -> None:
+        super().__init__(dataset, raise_on_error=raise_on_error)
+
+    def _compute(self) -> pd.DataFrame:
+        """Returns a tall DataFrame of all acquisition streams.
+
+        Includes a ``_stream_name`` column (e.g. ``"Behavior.HarpBehavior.PwmStart"``)
+        so the parquet can be filtered per stream downstream.
+        """
+        _ = self._dataset.load_all(strict=False)
+        frames = []
+        for stream in self._dataset.iter_all():
+            if stream.is_collection or stream.has_error:
+                if stream.is_collection:
+                    continue
+                logger.debug("Stream %s has error: %s", stream.name, stream.collect_errors())
+                if self._raise_on_error:
+                    raise ValueError(f"Stream {stream.name} has error")
+                continue
+            name = stream.resolved_name.replace("::", ".")
+            try:
+                df = stream.data.reset_index()
+                df["_stream_name"] = name
+                frames.append(df)
+            except Exception as exc:
+                if self._raise_on_error:
+                    raise
+                logger.debug("Could not load stream %s: %s", stream.name, exc)
+        return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()
+
+    def nwbize(self, nwb_file: ty.Any) -> ty.Any:
+        """Add all acquisition streams as DynamicTables to *nwb_file*."""
+        import contraqctor.contract as data_contract
+        import pynwb
+
+        from .helper import clean_dataframe_for_nwb
+
+        _ = self._dataset.load_all(strict=False)
+        for stream in self._dataset.iter_all():
+            if stream.is_collection:
+                err = stream.collect_errors()
+                if err:
+                    logger.debug("Collection stream %s has errors: %s", stream.name, err)
+                    if self._raise_on_error:
+                        raise ValueError(f"Collection stream {stream.name} has errors: {err}")
+                continue
+
+            name = stream.resolved_name.replace("::", ".")
+            try:
+                if stream.has_error:
+                    logger.debug("Stream %s has error: %s", stream.name, stream.collect_errors())
+                    if self._raise_on_error:
+                        raise ValueError(f"Stream {stream.name} has error: {stream.collect_errors()}")
+                    continue
+                if isinstance(stream, (data_contract.harp.HarpRegister, data_contract.csv.Csv)):
+                    table = pynwb.core.DynamicTable.from_dataframe(
+                        name=name, table_description=stream.description, df=stream.data.reset_index()
+                    )
+                    nwb_file.add_acquisition(table)
+                elif isinstance(stream, data_contract.json.SoftwareEvents):
+                    table = pynwb.core.DynamicTable.from_dataframe(
+                        name=name,
+                        table_description=stream.description,
+                        df=clean_dataframe_for_nwb(stream.data.reset_index()),
+                    )
+                    nwb_file.add_acquisition(table)
+                elif isinstance(stream, data_contract.json.PydanticModel):
+                    nwb_file.add_acquisition(
+                        pynwb.core.DynamicTable(
+                            name=name,
+                            description=stream.data.model_dump_json(),
+                        )
+                    )
+                else:
+                    raise ValueError(f"Stream {stream.name} has unsupported type {type(stream)}")
+            except Exception as exc:
+                if self._raise_on_error:
+                    raise
+                logger.debug("Error processing stream %s: %s", stream.name, exc)
+        return nwb_file

aind_behavior_vr_foraging_packaging-0.0.3/src/aind_behavior_vr_foraging_packaging/acquisition/helper.py

@@ -0,0 +1,85 @@
+import json
+from datetime import datetime
+from enum import Enum
+from typing import Any, Callable, Union
+
+import numpy as np
+import pandas as pd
+
+
+def clean_dataframe_for_nwb(data: pd.DataFrame) -> pd.DataFrame:
+    """
+    Clean a pandas DataFrame to ensure compatibility with NWB format.
+
+    Parameters
+    ----------
+    data : pd.DataFrame
+        The cleaned input DataFrame for NWB compatibility
+
+    Returns
+    -------
+    pd.DataFrame
+        A cleaned DataFrame that adheres to NWB data types
+    """
+    for column in data.columns:
+        # convert to nwb allowable types
+        data[column] = data[column].replace({None: np.nan})
+        data[column] = data[column].apply(lambda x: x.value if isinstance(x, Enum) else x)
+        data[column] = data[column].apply(lambda x: json.dumps(x) if isinstance(x, dict) else x)
+
+    return data
+
+
+_NestedStructureType = Union[dict, list, Any]
+
+
+def convert_values_in_nested_structure(
+    data: _NestedStructureType,
+    check_fn: Callable[[Any], bool],
+    convert_fn: Callable[[Any], Any],
+) -> _NestedStructureType:
+    """
+    Recursively convert values in nested dictionaries/lists based on a condition.
+
+    Parameters
+    ----------
+    data : _NestedStructureType
+        Input data structure which may contain nested dictionaries and lists.
+    check_fn : Callable
+        Function that returns True if value should be converted.
+    convert_fn : Callable
+        Function that converts the value.
+
+    Returns
+    -------
+    _NestedStructureType
+        Data structure with converted values.
+    """
+    if isinstance(data, dict):
+        return {k: convert_values_in_nested_structure(v, check_fn, convert_fn) for k, v in data.items()}
+    if isinstance(data, list):
+        return [convert_values_in_nested_structure(item, check_fn, convert_fn) for item in data]
+    return convert_fn(data) if check_fn(data) else data
+
+
+def convert_datetimes_to_iso_string(
+    data: _NestedStructureType,
+) -> _NestedStructureType:
+    """
+    Convert datetime objects in a nested structure to ISO format strings.
+
+    Parameters
+    ----------
+    data : _NestedStructureType
+        Input data structure which may contain nested dictionaries and lists.
+
+    Returns
+    -------
+    _NestedStructureType
+        Data structure with datetime objects converted to ISO format strings.
+    """
+    return convert_values_in_nested_structure(
+        data,
+        check_fn=lambda x: isinstance(x, datetime),
+        convert_fn=lambda x: x.isoformat(),
+    )

aind_behavior_vr_foraging_packaging-0.0.3/src/aind_behavior_vr_foraging_packaging/cli.py

@@ -0,0 +1,2 @@
+def main():
+    pass