ezmsg-xdf 0.1__py3-none-any.whl

ezmsg/xdf/__version__.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1'
+__version_tuple__ = version_tuple = (0, 1)

ezmsg/xdf/iter.py

@@ -0,0 +1,327 @@
+from dataclasses import replace
+from pathlib import Path
+import queue
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import pyxdf
+from ezmsg.lsl.util import AxisArray
+
+
+class XDFIterator:
+    def __init__(
+        self,
+        filepath: typing.Union[Path, str],
+        select: typing.Optional[
+            set[str]
+        ] = None,  # If set, then the iterator yields only AxisArray of selected stream(s).
+        # If None (default), then the iterator yields dicts with keys for each stream
+        chunk_dur: float = 1.0,  # Attempt to chunk data into chunks of this duration.
+        start_time: typing.Optional[float] = None,
+        stop_time: typing.Optional[float] = None,
+        rezero: bool = True,
+    ):
+        """
+        An Iterator that yields chunks from an XDF.
+        A typical offline analysis might load the entire file into memory, then perform a processing step on the entire
+        recording duration, and the next step on the entire result of the first step, and so on. This might require a
+        tremendous amount of memory and, if one is not careful about memory layout, can be incredibly slow. An
+        alternative procedure is to load the file into memory a chunk at a time (see Note1), then pass that chunk
+        through the entire processing pipeline, then proceed onto the next chunk (See Note2). We create an Iterator to
+        provide our chunks.
+        > Note1: I have not written a true lazy-loader for XDF because it has not yet been necessary as the files are
+          all small. Thus, I use pyxdf.load_xdf which loads the entire raw data into memory. The processing is still
+          done chunk-by-chunk.
+        > Note2: It should be possible to start on chunk[ix+1] while chunk[ix] is still going through the pipeline.
+          Indeed, this is (optionally) how it works online. However, the overhead of setting this up for offline
+          analysis is not worth the gain, at least not at this stage.
+
+        Args:
+            filepath: The path to the file to load and iterate over.
+            select: (Optional) A set of stream names to select. If None, then all streams are selected.
+            chunk_dur: The duration of each chunk in seconds.
+            start_time: Start playback at this time. If rezero is True then this is relative to the file start time.
+                If rezero is False then this is relative to the original timestamps.
+            stop_time: Truncate the playback to stop at this time. If rezero is True then this is relative to the file
+                start time. If rezero is False then this is relative to the original timestamps.
+            rezero: The absolute value of timestamps in an XDF file are useful for synchronization WITHIN file, but they
+                are absolutely meaningless outside the exact XDF file like in an ezmsg application. Thus, by default we
+                rezero the timestamps to start at t=0.0 for simplicity. However, there may be rare circumstances where
+                one wants to compare the timestamps produced by ezmsg to timestamps produced by another XDF analysis
+                tool that does not rezero. In that case, set rezero=False.
+        """
+        if isinstance(filepath, str):
+            filepath = Path(filepath).expanduser()
+        self._filepath = filepath
+        self._select = select
+        self._chunk_dur = chunk_dur
+        self._rezero = rezero
+        self._n_chunks = 0
+        self._t0 = 0.0
+        self._chunk_ix = 0
+        self._last_time = 0.0
+        self._metadata = {}
+        self._prev_file_read_s: float = (
+            0  # File read header in seconds for previous iteration
+        )
+        self._time_range: typing.Tuple[
+            typing.Optional[float], typing.Optional[float]
+        ] = (start_time, stop_time)
+        self._scan_file()
+
+    def _scan_file(self):
+        # Note: For larger datafiles we wouldn't want to load the entire thing into memory with load_xdf.
+        #  Instead, get a file handle, then
+        #  - Scan the file for chunk boundaries and timestamps
+        #  - Maintain a list of chunk boundaries
+        #  - Perform timestamp corrections (maintain corrected ts in memory or use func to correct during next pass?)
+        #  - Iterator operates on original chunk-boundaries, but using corrected timestamps.
+        #  However, we would need a custom file parser for that. For now, we load the relatively small
+        #  file into memory simply with pyxdf.load_xdf then iterate over the items in memory
+        #  at a user-defined chunk boundary (`chunk_dur`).
+        # Load xdf
+        self._streams, fileheader = pyxdf.load_xdf(
+            self._filepath,
+            select_streams=None
+            if (self._select is None or self._rezero)
+            else [{"name": _} for _ in self._select],
+        )
+        self._metadata = {}
+        self._file_read_s = 0
+        self._prev_file_read_s = 0
+        xdf_t0 = np.inf
+        xdf_tmax = 0
+        for strm in self._streams:
+            # Convert empty data to an array for easier slicing
+            if type(strm["time_series"]) is list:
+                strm["time_series"] = np.array(strm["time_series"])
+
+            # Get more digestable metadata
+            info = strm["info"]
+            new_meta = {
+                "name": info["name"][0],
+                "type": info["type"][0],
+                "channel_count": int(info["channel_count"][0]),
+                "nominal_srate": float(info["nominal_srate"][0]),
+            }
+            self._metadata[new_meta["name"]] = new_meta
+
+            # Update time range limits
+            tvec = strm["time_stamps"]
+            if len(tvec) > 0:
+                xdf_t0 = min(xdf_t0, tvec[0])
+                xdf_tmax = max(xdf_tmax, tvec[-1])
+
+        # Permanently modify streams' time stamps
+        if self._rezero:
+            for strm in self._streams:
+                strm["time_stamps"] = strm["time_stamps"] - xdf_t0
+            xdf_tmax -= xdf_t0
+            xdf_t0 = 0
+
+        # Adjust for provided time bounds
+        for strm in self._streams:
+            tvec = strm["time_stamps"]
+            if len(tvec) > 0:
+                b_keep = np.ones(len(tvec), dtype=bool)
+                if self._time_range[0] is not None:
+                    b_keep = np.logical_and(b_keep, tvec >= self._time_range[0])
+                if self._time_range[1] is not None:
+                    b_keep = np.logical_and(b_keep, tvec <= self._time_range[1])
+                if np.any(~b_keep):
+                    strm["time_stamps"] = tvec[b_keep]
+                    strm["timeseries"] = strm["timeseries"][b_keep]
+
+        # Recalculate tmax
+        xdf_dur = 0
+        for strm in self._streams:
+            tvec = strm["time_stamps"]
+            srate = float(strm["info"]["nominal_srate"][0])
+            adj = (1 / srate if srate > 0 else 0) - xdf_t0
+            if len(tvec) > 0:
+                xdf_dur = max(xdf_dur, tvec[-1] + adj)
+
+        # Chunking
+        self._n_chunks = int(np.ceil(xdf_dur / self._chunk_dur))
+        self._t0 = xdf_t0
+
+        # Drop streams that were not selected. (Could not drop earlier due to timestamp rezero)
+        if self._rezero and self._select is not None:
+            stream_names = [_["info"]["name"][0] for _ in self._streams]
+            self._streams = [self._streams[stream_names.index(_)] for _ in self._select]
+            self._metadata = {k: self._metadata[k] for k in self._select}
+
+        print(
+            f"Imported {len(self._streams)} streams from {self._filepath} "
+            f"spanning {xdf_dur:.2f} s beginning at t={xdf_t0:.2f}."
+        )
+
+    @property
+    def stream_meta(self) -> typing.Union[list[dict], dict]:
+        return self._metadata
+
+    @property
+    def n_chunks(self) -> int:
+        return self._n_chunks
+
+    def __iter__(self):
+        self._chunk_ix = 0
+        return self
+
+    def __next__(self) -> dict[str, tuple[npt.NDArray, npt.NDArray]]:
+        if self._chunk_ix >= self.n_chunks:
+            raise StopIteration
+        else:
+            out_dict = {}
+            t_start, t_stop = (
+                self._chunk_ix * self._chunk_dur + self._t0,
+                (self._chunk_ix + 1) * self._chunk_dur + self._t0,
+            )
+            for strm in self._streams:
+                b_chunk = np.logical_and(
+                    strm["time_stamps"] >= t_start, strm["time_stamps"] < t_stop
+                )
+                out_tvec = strm["time_stamps"][b_chunk]
+                out_data = strm["time_series"][b_chunk]
+                out_dict[strm["info"]["name"][0]] = (out_data, out_tvec)
+                if len(out_tvec) > 0:
+                    self._last_time = max(self._last_time, out_tvec[-1])
+            self._chunk_ix += 1
+            return out_dict
+
+
+def labels_from_strm(strm: dict) -> list[str]:
+    desc = strm["info"]["desc"][0]
+    if desc is not None and "channels" in desc:
+        labels = [_["label"][0] for _ in desc["channels"][0]["channel"]]
+    else:
+        n_ch = int(strm["info"]["channel_count"][0])
+        labels = [str(_ + 1) for _ in range(n_ch)]
+    return labels
+
+
+class XDFAxisArrayIterator(XDFIterator):
+    def __init__(self, *args, select: str, **kwargs):
+        """
+        This Iterator loads only a single stream and yields a single :obj:`AxisArray` object per chunk.
+
+        Args:
+            *args:
+            select: Unlike :obj:`XDFIterator`, this must be a single string, the name of the stream to select.
+            **kwargs:
+        """
+        kwargs["select"] = set((select,))
+        super().__init__(*args, **kwargs)
+        _sel = [_ for _ in self._select][0]
+        fs = self._metadata[_sel]["nominal_srate"] or 1.0
+        labels = labels_from_strm(self._streams[0])
+
+        self._template = AxisArray(
+            data=np.zeros(
+                (0, len(labels)), dtype=self._streams[0]["time_series"].dtype
+            ),
+            dims=["time", "ch"],
+            axes={
+                "time": AxisArray.Axis.TimeAxis(fs=fs, offset=0.0),
+                "ch": AxisArray.Axis.SpaceAxis(labels=labels),
+            },
+            key=self._streams[0]["info"]["name"][0],
+        )
+
+    def __next__(self) -> AxisArray:
+        result: typing.Optional[AxisArray] = None
+        chunk_dict = super().__next__()
+        # Should only be 1 in self._select. If there are more then we overwrite with the last.
+        for strm_name in self._select:
+            if strm_name in chunk_dict:
+                data, tvec = chunk_dict[strm_name]
+                result = replace(
+                    self._template,
+                    data=data,
+                    axes={
+                        **self._template.axes,
+                        "time": replace(
+                            self._template.axes["time"],
+                            offset=tvec[0] if len(tvec) else self._last_time,
+                        ),
+                    },
+                )
+        return result
+
+
+class XDFMultiAxArrIterator(XDFIterator):
+    def __init__(self, *args, force_single_sample: set = set(), **kwargs):
+        """
+        This Iterator loads multiple streams and yields a :obj:`AxisArray` object per iteration,
+        but the stream source might different between chunks.
+
+        Args:
+            *args:
+            force_single_sample: Use this to identify irregular-rate streams that might conceivably have more than one
+                event within the defined chunk_dur, for which :obj:`AxisArray` cannot represent timestamps properly.
+            **kwargs:
+        """
+        super().__init__(*args, **kwargs)
+        self._force_single_sample = force_single_sample
+        stream_names = [_["info"]["name"][0] for _ in self._streams]
+
+        # Create template messages for each stream
+        self._templates = {}
+        for stream_name, stream_meta in self._metadata.items():
+            stream = self._streams[stream_names.index(stream_name)]
+            labels = labels_from_strm(stream)
+            fs = stream_meta["nominal_srate"]
+            self._templates[stream_name] = AxisArray(
+                data=np.zeros(
+                    (0, stream_meta["channel_count"]), dtype=stream["time_series"].dtype
+                ),
+                dims=["time", "ch"],
+                axes={
+                    "time": AxisArray.Axis.TimeAxis(fs=fs or 1.0, offset=0.0),
+                    "ch": AxisArray.Axis.SpaceAxis(labels=labels),
+                },
+                key=stream_name,
+            )
+        self._pubqueue: queue.SimpleQueue[AxisArray] = queue.SimpleQueue()
+
+    def __next__(self) -> typing.Optional[AxisArray]:
+        if self._pubqueue.empty():
+            chunk_dict = super().__next__()
+            for k, template in self._templates.items():
+                if k in chunk_dict and len(chunk_dict[k][1]) > 0:
+                    data, tvec = chunk_dict[k]
+                    if k in self._force_single_sample:
+                        for ix, _t in enumerate(tvec):
+                            self._pubqueue.put_nowait(
+                                replace(
+                                    template,
+                                    data=data[ix : ix + 1],
+                                    axes={
+                                        **template.axes,
+                                        "time": replace(
+                                            template.axes["time"], offset=_t
+                                        ),
+                                    },
+                                )
+                            )
+                    else:
+                        self._pubqueue.put_nowait(
+                            replace(
+                                template,
+                                data=data,
+                                axes={
+                                    **template.axes,
+                                    "time": replace(
+                                        template.axes["time"],
+                                        offset=tvec[0]
+                                        if len(tvec)
+                                        else self._last_time,
+                                    ),
+                                },
+                            )
+                        )
+        try:
+            return self._pubqueue.get_nowait()
+        except queue.Empty:
+            return None

ezmsg/xdf/source.py

@@ -0,0 +1,109 @@
+import asyncio
+import os
+import typing
+from dataclasses import field
+
+import ezmsg.core as ez
+from ezmsg.util.generator import GenState
+from ezmsg.util.messages.axisarray import AxisArray
+
+from .iter import XDFAxisArrayIterator, XDFMultiAxArrIterator
+
+
+class XDFIteratorSettings(ez.Settings):
+    filepath: typing.Union[os.PathLike, str]
+    select: str
+    chunk_dur: float = 1.0
+    start_time: typing.Optional[float] = None
+    stop_time: typing.Optional[float] = None
+    rezero: bool = True
+    self_terminating: bool = False
+    """
+    If True, the unit will raise a :obj:`ez.NormalTermination` exception when the file is exhausted.
+    Note, however, that this will terminate the pipeline even if the data published by this unit are still in transit,
+    which will lead to the pipeline output being truncated before it has finished processing the stream.
+    `self_terminating` should only be used when it is not important that the pipeline finish processing data, such
+    as during prototyping and testing. 
+    """
+
+
+class XDFIteratorUnit(ez.Unit):
+    STATE = GenState
+    SETTINGS = XDFIteratorSettings
+
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    OUTPUT_TERM = ez.OutputStream(typing.Any)
+
+    def initialize(self) -> None:
+        self.construct_generator()
+
+    def construct_generator(self):
+        self.STATE.gen = XDFAxisArrayIterator(
+            filepath=self.SETTINGS.filepath,
+            select=self.SETTINGS.select,
+            chunk_dur=self.SETTINGS.chunk_dur,
+            start_time=self.SETTINGS.start_time,
+            stop_time=self.SETTINGS.stop_time,
+            rezero=self.SETTINGS.rezero,
+        )
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def pub_chunk(self) -> typing.AsyncGenerator:
+        try:
+            while True:
+                msg = next(self.STATE.gen)
+                if msg.data.size > 0:
+                    yield self.OUTPUT_SIGNAL, msg
+                else:
+                    await asyncio.sleep(0)
+        except StopIteration:
+            ez.logger.debug(
+                f"File ({self.SETTINGS.filepath} :: {self.SETTINGS.select}) exhausted."
+            )
+            if self.SETTINGS.self_terminating:
+                raise ez.NormalTermination
+            yield self.OUTPUT_TERM, True
+
+
+class XDFMultiIteratorUnitSettings(XDFIteratorSettings):
+    select: typing.Optional[set[str]] = None  # Override with a default
+    force_single_sample: set = field(default_factory=set)
+
+
+class XDFMultiIteratorUnit(ez.Unit):
+    STATE = GenState
+    SETTINGS = XDFMultiIteratorUnitSettings
+
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    OUTPUT_TERM = ez.OutputStream(typing.Any)
+
+    def initialize(self) -> None:
+        self.construct_generator()
+
+    def construct_generator(self):
+        self.STATE.gen = XDFMultiAxArrIterator(
+            filepath=self.SETTINGS.filepath,
+            select=self.SETTINGS.select,
+            chunk_dur=self.SETTINGS.chunk_dur,
+            start_time=self.SETTINGS.start_time,
+            stop_time=self.SETTINGS.stop_time,
+            rezero=self.SETTINGS.rezero,
+            force_single_sample=self.SETTINGS.force_single_sample,
+        )
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def pub_multi(self) -> typing.AsyncGenerator:
+        try:
+            while True:
+                msg = next(self.STATE.gen)
+                if msg is not None:
+                    yield self.OUTPUT_SIGNAL, msg
+                else:
+                    await asyncio.sleep(0)
+        except StopIteration:
+            ez.logger.debug(
+                f"File ({self.SETTINGS.filepath} :: {self.SETTINGS.select}) exhausted."
+            )
+            if self.SETTINGS.self_terminating:
+                raise ez.NormalTermination
+            yield self.OUTPUT_TERM, True

ezmsg_xdf-0.1.dist-info/METADATA

@@ -0,0 +1,18 @@
+Metadata-Version: 2.3
+Name: ezmsg-xdf
+Version: 0.1
+Summary: Namespace package for ezmsg to load iterate data from xdf files
+Author-email: Chadwick Boulay <chadwick.boulay@gmail.com>
+License-File: LICENSE
+Requires-Python: >=3.9
+Requires-Dist: ezmsg-lsl>=0.5.0
+Requires-Dist: ezmsg>=3.5.0
+Requires-Dist: numpy>=2.0.2
+Requires-Dist: pyxdf>=1.16.8
+Provides-Extra: test
+Requires-Dist: pytest>=8.3.3; extra == 'test'
+Description-Content-Type: text/markdown
+
+# ezmsg.xdf
+
+ezmsg namespace package for working with XDF files.

ezmsg_xdf-0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+ezmsg/xdf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/xdf/__version__.py,sha256=7V2ukxSfzXlOjRj21NN9XMWrUwTZJIQZggSzFRx5qs8,406
+ezmsg/xdf/iter.py,sha256=7a4oJNIHoGJAobxQRl3iMYtRn-dh_XtxDMqdlqQTsy0,14396
+ezmsg/xdf/source.py,sha256=jAG6WacFJWCsO8qConswwiQKEWcO50QLBZWGNJUJ_hI,3760
+ezmsg_xdf-0.1.dist-info/METADATA,sha256=WxR8WM5xXjE4nnT6xd4FV0LEdDr-CF0CEougHgoNURg,517
+ezmsg_xdf-0.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+ezmsg_xdf-0.1.dist-info/licenses/LICENSE,sha256=KImy0vSiKLkW28zWqI8rIaLwvjGKxOjE5kZmVxZDdy4,1066
+ezmsg_xdf-0.1.dist-info/RECORD,,

ezmsg_xdf-0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.25.0
+Root-Is-Purelib: true
+Tag: py3-none-any

ezmsg_xdf-0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 ezmsg-org
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.