Source code for vivyd.data.signal

from __future__ import annotations

from quantities import UnitQuantity, Quantity

from dataclasses import dataclass
from numpy.typing import NDArray
from numpy import asarray
from uuid import uuid4
from zarr import Group, open_array
from zarr.storage import ZipStore
from typing import Any




@dataclass(kw_only=True)
class Signal:
    """
    A class representing a signal, which is a time series of data points.
    
    Parameters
    ----------
    name : str
        The name of the signal.
    time : NDArray
        The time points of the signal.
    data : NDArray
        The data points of the signal. The first dimension should correspond
        to time, and the remaining dimensions can be used for multi-dimensional
        signals (e.g., multi-channel recordings).
    units : UnitQuantity | Quantity
        The units of the value, as defined by the ``quantities`` library.
    """
    
    name:  str
    data:  NDArray
    units: UnitQuantity | Quantity
    time:  NDArray
    
    def __post_init__(self):
        if self.time.ndim != 1:
            raise ValueError(f"Time array must be 1-dimensional. Got {self.time.ndim} dimensions.")
        if self.data.shape[0] != len(self.time):
            raise ValueError(f"Data and time arrays must have the same length. Got {self.data.shape[0]} and {len(self.time)}.")
    
    def _save_arrays(self, root: Group) -> tuple[str, str]:
        """Save arrays to Zarr and return their paths."""
        uid = str(uuid4())

        data_path = f"arrays/{uid}/data"
        time_path = f"arrays/{uid}/time"

        # Create zarr arrays via group methods with chunking
        # Convert to numpy in case they're already zarr arrays
        data_arr = asarray(self.data)
        time_arr = asarray(self.time)
        
        root.create_array(
            data_path,
            data   = data_arr,
            chunks = (min(data_arr.shape[0], 1_000_000),),
        )

        root.create_array(
            time_path,
            data   = time_arr,
            chunks = (min(time_arr.shape[0], 1_000_000),),
        )

        return data_path, time_path
    
    @property
    def __dict__(self) -> dict[str, Any]:
        return {
            "name" : self.name,
            "type" : self.__class__.__name__,
            "units": self.units.dimensionality.string,
        }
        
    @staticmethod
    def _from_dict(data: dict[str, Any], store: ZipStore) -> Signal:
        data_zarr = open_array(store, path=data["data_ref"], mode="r")
        time_zarr = open_array(store, path=data["time_ref"], mode="r")
        return Signal(
            name  = data["name"],
            units = Quantity(1.0, data["units"]).units,
            data  = asarray(data_zarr),
            time  = asarray(time_zarr),
        )
        


    def tree(self, indent: int = 2) -> str:
        """
        Parameters
        ----------
        indent : int, optional
            The number of spaces to indent the tree representation of the signal.
            Default is 2.
        
        Returns
        -------
        str
            The tree representation of the signal.
        """
        
        type_  = self.__class__.__name__
        units_ = self.units.dimensionality.string
        shape_ = self.data.shape
        return f"{' ' * indent}<{type_}> {self.name} [{units_}] {shape_}"