from __future__ import annotations
from .container import Container
from .study import Study
from .signal import Signal
import json
import zarr
import warnings
from zarr.storage import ZipStore
from dataclasses import dataclass, field
from typing import Any
from pathlib import Path
from zenodo_get import download
[docs]
@dataclass(kw_only=True)
class DataSet(Container):
"""
A class representing a dataset, which is a hierarchical collection of
studies. It is the top-level container in the data organization structure.
A :class:`DataSet` is meant to hold an entire experimental campaign.
Parameters
----------
name : str
The name of the dataset.
site : Site
Information about the site/location where the data was collected.
authors : list[Author]
The authors of the dataset.
info : dict[str, Any]
Additional metadata about the dataset, such as experimental conditions,
measurement techniques, etc. This is a flexible field that can hold any
relevant information that does not fit into the predefined fields.
"""
site : Site
"""Information about the site/location where the data was collected."""
authors: list[Author] = field(default_factory=list)
"""The authors of the dataset."""
[docs]
def add(self, *study: Study) -> DataSet:
"""
Add one or more studies to the dataset.
Parameters
----------
*study : Study
One or more study objects to add to the dataset.
Returns
-------
DataSet
The dataset with the added studies.
"""
for s in study:
super().add(s.name, s)
return self
[docs]
def save(self, path: Path | str, create_dirs: bool = False):
"""
Save the dataset to a file. The dataset is saved in a chunked compressed
format using Zarr, which allows for efficient storage and retrieval of
large datasets. The metadata (structure and references to data) is
stored as a JSON string in the root attributes of the Zarr store.
Parameters
----------
path : Path | str
The file path where the dataset should be saved. Must have a
``.vxp`` extension.
create_dirs : bool, optional
If `True`, create any missing directories in the specified path.
Default is `False`.
Raises
------
ValueError
If the provided path does not have a ``.vxp`` extension.
"""
if isinstance(path, str):
path = Path(path)
if path.suffix != ".vxp":
raise ValueError("Use .vxp for chunked storage")
if create_dirs:
path.parent.mkdir(parents=True, exist_ok=True)
# Suppress zarr/zipfile duplicate name warning (harmless internal behavior)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning, message=".*Duplicate name.*")
store = ZipStore(str(path), mode="w") # Zarr store backed by a zip file
root = zarr.open_group(store=store, mode="w") # Root of the Zarr store
# Walk the tree and save arrays, collecting refs
def inject_refs(obj):
# Recursively store contents of containers
if isinstance(obj, Container):
return {
**obj.__dict__,
"contents": [inject_refs(c) for c in obj.contents.values()]
}
# For signals, compress arrays (heavy data) to seperate zarr files
# Only references to arrays are stored in the metadata
elif isinstance(obj, Signal):
data_ref, time_ref = obj._save_arrays(root) # Heavy signal data is saved in a compressed, chunked format
return {
**obj.__dict__,
"data_ref": data_ref,
"time_ref": time_ref
}
else:
return obj.__dict__
meta = inject_refs(self) # Apply to the entire dataset
# Store metadata as JSON string in the store's root attributes
root.attrs["_metadata"] = json.dumps(meta)
store.close()
[docs]
@staticmethod
def load(path: Path | str) -> DataSet:
"""
Parameters
----------
path : Path | str
The file path from which to load the dataset, which must respect the
``.vxp`` extension format used for chunked storage.
Returns
-------
DataSet
The loaded dataset.
Raises
------
ValueError
If the provided path does not have a ``.vxp`` extension.
"""
if isinstance(path, str):
path = Path(path)
if path.suffix != ".vxp":
raise ValueError("Use .vxp for chunked storage")
store = ZipStore(str(path), mode="r")
root = zarr.open_group(store=store, mode="r")
# Read metadata from zarr group attributes (stored as JSON string)
meta_str = str(root.attrs["_metadata"])
meta = json.loads(meta_str)
return DataSet._from_dict(meta, store)
[docs]
@staticmethod
def download_from_zenodo(
record_id: str,
output_dir: Path | str,
token: str | None = None,
create_dirs: bool = False
):
"""
Download a dataset from Zenodo using the record ID. The dataset is
expected to be in a valid ``.vxp`` format and will be saved to the
specified output directory.
Parameters
----------
record_id : str
The Zenodo record ID of the dataset to download.
output_dir : Path | str
The directory where the downloaded dataset should be saved.
token : str | None, optional
An optional Zenodo access token for downloading private records. If
the record is public, this can be left as `None`. Default is `None`.
create_dirs : bool, optional
If `True`, create any missing directories in the specified output
path. Default is `False`.
Tip
---
You can find the record ID in the URL of the Zenodo page for the
dataset. For example, if the URL is
`https://zenodo.org/record/12345678`, the record ID is `12345678`.
"""
if isinstance(output_dir, str):
output_dir = Path(output_dir)
if create_dirs:
output_dir.mkdir(parents=True, exist_ok=True)
download(
record=record_id,
access_token=token,
output_dir=output_dir
)
@property
def __dict__(self) -> dict[str, Any]:
return {
**super().__dict__,
"site" : self.site.__dict__,
"authors": [author.__dict__ for author in self.authors],
}
@staticmethod
def _from_dict(data: dict[str, Any], store: ZipStore) -> DataSet:
new = DataSet(
name = data["name"],
info = data["info"],
site = Site(**data["site"]),
authors = [Author(**a) for a in data["authors"]]
)
for study in data["contents"]:
new.add(Study._from_dict(study, store))
return new
[docs]
@dataclass(kw_only=True)
class Author:
"""
A class representing an author of a dataset. This is used to store
information about the individuals who contributed to the creation of the
dataset.
Parameters
----------
info : dict[str, Any]
Additional metadata about the author, such as their role in the project,
contact information, etc. This is a flexible field that can hold any
relevant information that does not fit into the predefined fields.
"""
firstname : str
lastname : str
email : str
institution: str
info : dict[str, Any] = field(default_factory=dict)
"""
Additional metadata about the author, such as their role in the project,
contact information, etc. This is a flexible field that can hold any
relevant information that does not fit into the predefined fields.
"""
[docs]
@dataclass(kw_only=True)
class Site:
"""
A class representing the site/location where the data was collected. This
can be a lab or a field site for instance.
"""
name : str
location : str
institution: str