Merge pull request #15 from ezmsg-org/dev

CoordinateAxis for irregular rate streams

Author: cboulay

.github/workflows/python-tests.yml

@@ -11,7 +11,7 @@ jobs:
   build:
     strategy:
       matrix:
-        python-version: [3.9, "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
         os:
           - "ubuntu-latest"
           - "windows-latest"
@@ -21,6 +21,32 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - name: Install linux dependencies
+        if: ${{ runner.os == 'Linux' }}
+        shell: bash
+        run: |
+          sudo apt update
+          sudo apt install -y binutils libpugixml-dev qtbase5-dev qt5-qmake
+
+      - name: Install liblsl (linux)
+        if: ${{ runner.os == 'Linux' }}
+        shell: bash
+        run: |
+          curl -L https://github.com/sccn/liblsl/releases/download/v1.16.2/liblsl-1.16.2-jammy_amd64.deb -o liblsl-1.16.2-jammy_amd64.deb
+          sudo apt install -y ./liblsl-1.16.2-jammy_amd64.deb
+          rm liblsl-1.16.2-jammy_amd64.deb
+
+      - name: Install liblsl (macOS)
+        if: ${{ runner.os == 'macOS' }}
+        shell: bash
+        run: |
+          curl -L https://github.com/sccn/liblsl/releases/download/v1.16.0/liblsl-1.16.0-OSX_arm64.tar.bz2 -o liblsl-1.16.0-OSX_arm64.tar.bz2
+          tar -xf liblsl-1.16.0-OSX_arm64.tar.bz2
+          mv lib/liblsl.1.16.0.dylib .
+          echo "PYLSL_LIB=$PWD/liblsl.1.16.0.dylib" >> $GITHUB_ENV
+          rm -R lib include bin
+          rm liblsl-1.16.0-OSX_arm64.tar.bz2
+
       - name: Install uv
         uses: astral-sh/setup-uv@v2
         with:
@@ -33,9 +59,10 @@ jobs:
       - name: Install the project
         run: uv sync --all-extras --dev
 
-      - name: Lint
-        run:
-          uv tool run ruff check --output-format=github src
+      - name: Ruff check
+        uses: astral-sh/ruff-action@v1
+        with:
+          src: "./src"
 
-#      - name: Run tests
-#        run: uv run pytest tests
+      - name: Run tests
+        run: uv run pytest tests

examples/lsl_inlet_example.py

@@ -1,8 +1,9 @@
 from typing import Optional
 
 import ezmsg.core as ez
-from ezmsg.lsl.units import LSLInletUnit, LSLInletSettings
+from ezmsg.lsl.units import LSLInletUnit, LSLInletSettings, LSLInfo
 from ezmsg.util.debuglog import DebugLog, DebugLogSettings
+import typer
 # from ezmsg.util.messagelogger import MessageLogger, MessageLoggerSettings
 
 
@@ -23,8 +24,10 @@ class LSLDemoSystem(ez.Collection):
     def configure(self) -> None:
         self.INLET.apply_settings(
             LSLInletSettings(
-                stream_name=self.SETTINGS.stream_name,
-                stream_type=self.SETTINGS.stream_type,
+                info=LSLInfo(
+                    name=self.SETTINGS.stream_name,
+                    type=self.SETTINGS.stream_type,
+                )
             )
         )
         self.LOGGER.apply_settings(
@@ -38,8 +41,13 @@ def network(self) -> ez.NetworkDefinition:
         return ((self.INLET.OUTPUT_SIGNAL, self.LOGGER.INPUT),)
 
 
-if __name__ == "__main__":
-    # Run the websocket system
+def main(stream_name: str = "", stream_type: str = "EEG"):
     system = LSLDemoSystem()
-    system.apply_settings(LSLDemoSystemSettings(stream_name="", stream_type="EEG"))
+    system.apply_settings(
+        LSLDemoSystemSettings(stream_name=stream_name, stream_type=stream_type)
+    )
     ez.run(SYSTEM=system)
+
+
+if __name__ == "__main__":
+    typer.run(main)

examples/lsl_outlet_example.py

@@ -8,6 +8,7 @@
 import ezmsg.core as ez
 from ezmsg.sigproc.synth import EEGSynth, EEGSynthSettings
 from ezmsg.lsl.units import LSLOutletUnit, LSLOutletSettings
+import typer
 
 
 class LSLDemoSystemSettings(ez.Settings):
@@ -48,10 +49,14 @@ def process_components(self) -> Tuple[ez.Component, ...]:
         return self.EEG, self.OUTLET
 
 
-if __name__ == "__main__":
+def main(stream_name: str = "ezmsg-EEGSynth", stream_type: str = "EEG"):
     # Run the websocket system
     system = LSLDemoSystem()
     system.apply_settings(
-        LSLDemoSystemSettings(stream_name="ezmsg-EEGSynth", stream_type="EEG")
+        LSLDemoSystemSettings(stream_name=stream_name, stream_type=stream_type)
     )
     ez.run(SYSTEM=system)
+
+
+if __name__ == "__main__":
+    typer.run(main)

pyproject.toml

@@ -7,15 +7,15 @@ authors = [
 readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [
-    "ezmsg>=3.5.0",
+    "ezmsg>=3.6.0",
     "numpy>=1.26.4",
     "pylsl>=1.16.2",
 ]
 dynamic = ["version"]
 
 [project.optional-dependencies]
 test = [
-    "ezmsg-sigproc>=1.3.2",
+    "ezmsg-sigproc>=1.5.0",
     "flake8>=7.1.1",
     "pytest-cov>=5.0.0",
     "pytest>=8.3.3",
@@ -37,6 +37,7 @@ packages = ["src/ezmsg"]
 [tool.uv]
 dev-dependencies = [
     "ruff>=0.6.6",
+    "typer>=0.13.0",
 ]
 
 [tool.uv.sources]

src/ezmsg/lsl/inlet.py

@@ -1,14 +1,16 @@
 import asyncio
-from dataclasses import dataclass, field, fields, replace
+from dataclasses import dataclass, field, fields
 import time
 import typing
 
 import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 import numpy as np
 import numpy.typing as npt
 import pylsl
 
-from .util import AxisArray
+from .util import ClockSync
 
 
 fmt2npdtype = {
@@ -76,42 +78,6 @@ class LSLInletState(ez.State):
     clock_offset: float = 0.0
 
 
-class ClockSync:
-    def __init__(self, alpha: float = 0.1, min_interval: float = 0.5):
-        self.alpha = alpha
-        self.min_interval = min_interval
-
-        self.offset = 0.0
-        self.last_update = 0.0
-        self.count = 0
-
-    async def update(self, force: bool = False, burst: int = 4) -> None:
-        dur_since_last = time.time() - self.last_update
-        dur_until_next = self.min_interval - dur_since_last
-        if force or dur_until_next <= 0:
-            offsets = []
-            for _ in range(burst):
-                if self.count % 2:
-                    y, x = time.time(), pylsl.local_clock()
-                else:
-                    x, y = pylsl.local_clock(), time.time()
-                offsets.append(y - x)
-                self.last_update = y
-                await asyncio.sleep(0.001)
-            offset = np.mean(offsets)
-
-            if self.count > 0:
-                # Exponential decay smoothing
-                offset = (1 - self.alpha) * self.offset + self.alpha * offset
-            self.offset = offset
-            self.count += burst
-        else:
-            await asyncio.sleep(dur_until_next)
-
-    def convert_timestamp(self, lsl_timestamp: float) -> float:
-        return lsl_timestamp + self.offset
-
-
 class LSLInletUnit(ez.Unit):
     """
     Represents a node in a graph that creates an LSL inlet and
@@ -223,22 +189,29 @@ def _reset_inlet(self) -> None:
                 ch_labels.append(str(len(ch_labels) + 1))
             # Pre-allocate a message template.
             fs = inlet_info.nominal_srate()
+            time_ax = (
+                AxisArray.TimeAxis(fs=fs)
+                if fs
+                else AxisArray.CoordinateAxis(
+                    data=np.array([]), dims=["time"], unit="s"
+                )
+            )
             self._msg_template = AxisArray(
                 data=np.empty((0, n_ch)),
                 dims=["time", "ch"],
                 axes={
-                    "time": AxisArray.Axis.TimeAxis(
-                        fs=fs if fs else 1.0
-                    ),  # HACK: Use 1.0 for irregular rate.
-                    "ch": AxisArray.Axis.SpaceAxis(labels=ch_labels),
+                    "time": time_ax,
+                    "ch": AxisArray.CoordinateAxis(
+                        data=np.array(ch_labels), dims=["ch"]
+                    ),
                 },
                 key=inlet_info.name(),
             )
 
     async def initialize(self) -> None:
         self._reset_resolver()
         self._reset_inlet()
-        # TODO: Let the clock_sync task do its job at the beginning.
+        await self.clock_sync.update(force=True, burst=1000)
 
     def shutdown(self) -> None:
         if self.STATE.inlet is not None:
@@ -252,8 +225,7 @@ def shutdown(self) -> None:
     @ez.task
     async def clock_sync_task(self) -> None:
         while True:
-            force = self.clock_sync.count < 1000
-            await self.clock_sync.update(force=force, burst=1000 if force else 4)
+            await self.clock_sync.update(force=False, burst=4)
 
     @ez.subscriber(INPUT_SETTINGS)
     async def on_settings(self, msg: LSLInletSettings) -> None:
@@ -292,36 +264,32 @@ async def lsl_pull(self) -> typing.AsyncGenerator:
                     if samples is None
                     else samples
                 )
+
                 if self.SETTINGS.use_arrival_time:
-                    # time.time() gives us NOW, but we want the timestamp of the 0th sample in the chunk
-                    t0 = time.time() - (timestamps[-1] - timestamps[0])
+                    timestamps = time.time() - (timestamps - timestamps[0])
                 else:
-                    t0 = self.clock_sync.convert_timestamp(timestamps[0])
+                    timestamps = self.clock_sync.lsl2system(timestamps)
+
                 if self.SETTINGS.info.nominal_srate <= 0.0:
-                    # Irregular rate streams need to be streamed sample-by-sample
-                    for ts, samp in zip(timestamps, data):
-                        out_msg = replace(
-                            self._msg_template,
-                            data=samp[None, ...],
-                            axes={
-                                **self._msg_template.axes,
-                                "time": replace(
-                                    self._msg_template.axes["time"],
-                                    offset=t0 + (ts - timestamps[0]),
-                                ),
-                            },
-                        )
-                        yield self.OUTPUT_SIGNAL, out_msg
+                    # Irregular rate stream uses CoordinateAxis for time so each sample has a timestamp.
+                    out_time_ax = replace(
+                        self._msg_template.axes["time"],
+                        data=np.array(timestamps),
+                    )
                 else:
-                    # Regular-rate streams can go in a chunk
-                    out_msg = replace(
-                        self._msg_template,
-                        data=data,
-                        axes={
-                            **self._msg_template.axes,
-                            "time": replace(self._msg_template.axes["time"], offset=t0),
-                        },
+                    # Regular rate uses a LinearAxis for time so we only need the time of the first sample.
+                    out_time_ax = replace(
+                        self._msg_template.axes["time"], offset=timestamps[0]
                     )
-                    yield self.OUTPUT_SIGNAL, out_msg
+
+                out_msg = replace(
+                    self._msg_template,
+                    data=data,
+                    axes={
+                        **self._msg_template.axes,
+                        "time": out_time_ax,
+                    },
+                )
+                yield self.OUTPUT_SIGNAL, out_msg
             else:
                 await asyncio.sleep(0.001)