BulletLab 0.1.2__py3-none-any.whl

bulletlab/logging/json_writer.py

@@ -0,0 +1,70 @@
+"""
+JSON/NDJSON writer backend for DataLogger.
+
+Internal module — used by :class:`~bulletlab.logging.logger.DataLogger`.
+Writes newline-delimited JSON (one JSON object per line) for streaming-friendly
+log files.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+
+class JsonWriter:
+    """Writes log data as newline-delimited JSON (NDJSON format).
+
+    Each call to :meth:`write_row` appends a single JSON object to the file.
+
+    Args:
+        filepath: Output file path.
+
+    Example::
+
+        writer = JsonWriter("run1.json")
+        writer.write_row({"t": 0.0, "speed": 1.5, "roll": 0.01})
+        writer.close()
+    """
+
+    def __init__(self, filepath: str | Path) -> None:
+        self._filepath = Path(filepath)
+        self._file = open(self._filepath, "w", encoding="utf-8")
+
+    def write_row(self, row: dict[str, Any]) -> None:
+        """Write a single record as a JSON line.
+
+        Args:
+            row: Dictionary of field names to values.
+        """
+        self._file.write(json.dumps(row, default=_json_default) + "\n")
+
+    def flush(self) -> None:
+        """Flush the write buffer to disk."""
+        self._file.flush()
+
+    def close(self) -> None:
+        """Close the file handle."""
+        if not self._file.closed:
+            self._file.flush()
+            self._file.close()
+
+    def __del__(self) -> None:
+        self.close()
+
+
+def _json_default(obj: Any) -> Any:
+    """Fallback JSON serializer for numpy types and similar."""
+    try:
+        import numpy as np  # noqa: PLC0415
+
+        if isinstance(obj, (np.integer,)):
+            return int(obj)
+        if isinstance(obj, (np.floating,)):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+    except ImportError:
+        pass
+    return str(obj)

bulletlab/logging/logger.py

@@ -0,0 +1,258 @@
+"""
+DataLogger – records time-series experiment data to CSV or JSON files.
+
+The DataLogger watches callable data sources and writes them to a file on
+every call to :meth:`step`. Supports CSV and newline-delimited JSON (NDJSON)
+formats, determined by the file extension.
+
+Example::
+
+    from bulletlab.logging import DataLogger
+
+    logger = DataLogger()
+    logger.watch("speed",  lambda: robot.speed)
+    logger.watch("roll",   lambda: robot.roll)
+    logger.watch("height", lambda: robot.base_position[2])
+    logger.start("experiment_01.csv")
+
+    for _ in range(1000):
+        sim.step()
+        logger.step()
+
+    logger.stop()
+
+Context manager usage::
+
+    with DataLogger() as logger:
+        logger.watch("speed", lambda: robot.speed)
+        logger.start("run.csv")
+        for _ in range(1000):
+            sim.step()
+            logger.step()
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Any, Callable
+
+from bulletlab.logging.csv_writer import CsvWriter
+from bulletlab.logging.json_writer import JsonWriter
+
+
+class DataLogger:
+    """Records experiment data from callable sources to CSV or JSON files.
+
+    Channels can be registered before or after calling :meth:`start`.
+    The logger automatically determines the file format from the extension:
+    ``.csv`` → CSV, ``.json`` or ``.ndjson`` → NDJSON.
+
+    Args:
+        include_timestamp: If ``True``, prepend a ``"t"`` column with the
+            wall-clock time of each step. Defaults to ``True``.
+        include_step: If ``True``, prepend a ``"step"`` column with the step
+            index. Defaults to ``True``.
+
+    Example::
+
+        logger = DataLogger()
+        logger.watch("speed", lambda: robot.speed)
+        logger.start("run1.csv")
+        for _ in range(1000):
+            sim.step()
+            logger.step()
+        logger.stop()
+    """
+
+    def __init__(
+        self,
+        include_timestamp: bool = True,
+        include_step: bool = True,
+    ) -> None:
+        self._channels: dict[str, Callable[[], Any]] = {}
+        self._writer: CsvWriter | JsonWriter | None = None
+        self._step_count: int = 0
+        self._include_timestamp = include_timestamp
+        self._include_step = include_step
+        self._start_time: float = 0.0
+        self._running: bool = False
+
+    # ------------------------------------------------------------------
+    # Channel registration
+    # ------------------------------------------------------------------
+
+    def watch(self, name: str, source: Callable[[], Any]) -> "DataLogger":
+        """Register a data source to log.
+
+        Args:
+            name: Column name in the output file.
+            source: Callable returning the current value.
+
+        Returns:
+            self, for method chaining.
+
+        Example::
+
+            logger.watch("speed", lambda: robot.speed)
+            logger.watch("roll",  lambda: robot.roll)
+        """
+        self._channels[name] = source
+        return self
+
+    def unwatch(self, name: str) -> None:
+        """Remove a previously registered channel.
+
+        Args:
+            name: Channel name to remove.
+        """
+        self._channels.pop(name, None)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(self, filepath: str | Path) -> "DataLogger":
+        """Open the output file and begin logging.
+
+        Args:
+            filepath: Path to the output file. Extension determines format:
+                ``.csv`` for CSV, ``.json`` / ``.ndjson`` for NDJSON.
+
+        Returns:
+            self, for method chaining.
+
+        Raises:
+            RuntimeError: If the logger is already running.
+
+        Example::
+
+            logger.start("run1.csv")
+        """
+        if self._running:
+            raise RuntimeError("DataLogger is already running. Call stop() first.")
+
+        filepath = Path(filepath)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        ext = filepath.suffix.lower()
+
+        fieldnames = self._build_fieldnames()
+
+        if ext in (".json", ".ndjson", ".jsonl"):
+            self._writer = JsonWriter(filepath)
+        else:
+            # Default to CSV
+            self._writer = CsvWriter(filepath, fieldnames)
+
+        self._step_count = 0
+        self._start_time = time.monotonic()
+        self._running = True
+        return self
+
+    def stop(self) -> None:
+        """Flush and close the output file.
+
+        Example::
+
+            logger.stop()
+        """
+        if self._writer is not None:
+            self._writer.close()
+            self._writer = None
+        self._running = False
+
+    # ------------------------------------------------------------------
+    # Data recording
+    # ------------------------------------------------------------------
+
+    def step(self, t: float | None = None) -> dict[str, Any]:
+        """Record one row of data from all watched sources.
+
+        Call this once per simulation step. If not started, this is a no-op.
+
+        Args:
+            t: Timestamp override. If ``None``, uses wall-clock elapsed time
+                since :meth:`start` was called.
+
+        Returns:
+            The recorded row as a dictionary.
+
+        Example::
+
+            for _ in range(1000):
+                sim.step()
+                logger.step()
+        """
+        if not self._running or self._writer is None:
+            return {}
+
+        timestamp = t if t is not None else (time.monotonic() - self._start_time)
+        row: dict[str, Any] = {}
+
+        if self._include_step:
+            row["step"] = self._step_count
+        if self._include_timestamp:
+            row["t"] = round(timestamp, 6)
+
+        for name, source in self._channels.items():
+            try:
+                val = source()
+            except Exception:
+                val = float("nan")
+            row[name] = val
+
+        self._writer.write_row(row)
+        self._step_count += 1
+        return row
+
+    def flush(self) -> None:
+        """Manually flush buffered data to disk."""
+        if self._writer is not None:
+            self._writer.flush()
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _build_fieldnames(self) -> list[str]:
+        """Build the ordered list of field names including meta-columns."""
+        fields: list[str] = []
+        if self._include_step:
+            fields.append("step")
+        if self._include_timestamp:
+            fields.append("t")
+        fields.extend(self._channels.keys())
+        return fields
+
+    # ------------------------------------------------------------------
+    # State
+    # ------------------------------------------------------------------
+
+    @property
+    def is_running(self) -> bool:
+        """``True`` if the logger is currently recording."""
+        return self._running
+
+    @property
+    def step_count(self) -> int:
+        """Number of rows recorded since :meth:`start` was last called."""
+        return self._step_count
+
+    @property
+    def channel_names(self) -> list[str]:
+        """Names of all registered channels."""
+        return list(self._channels.keys())
+
+    # ------------------------------------------------------------------
+    # Context manager
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> "DataLogger":
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self.stop()
+
+    def __repr__(self) -> str:
+        status = "running" if self._running else "stopped"
+        return f"DataLogger({status}, channels={list(self._channels.keys())}, steps={self._step_count})"

bulletlab/plotting/__init__.py

@@ -0,0 +1,9 @@
+"""
+BulletLab plotting subpackage.
+
+Provides LivePlot for real-time data visualization using PyQtGraph.
+"""
+
+from bulletlab.plotting.live_plot import LivePlot
+
+__all__ = ["LivePlot"]

bulletlab/plotting/live_plot.py

@@ -0,0 +1,364 @@
+"""
+LivePlot – real-time data visualization using PyQtGraph.
+
+LivePlot opens a PyQtGraph window and plots live data from callable sources.
+Multiple traces can be added with custom colors. The plot supports zoom,
+pan (via PyQtGraph's native interaction), pause, resume, and image export.
+
+Example::
+
+    from bulletlab.plotting import LivePlot
+
+    plot = LivePlot(title="Robot Telemetry", max_points=500)
+    plot.watch("Speed",  lambda: robot.speed, color="#00ff88")
+    plot.watch("Roll",   lambda: robot.roll,  color="#ff4488")
+    plot.watch("Height", lambda: robot.base_position[2], color="#44aaff")
+    plot.start()
+
+    for _ in range(5000):
+        sim.step()
+        plot.update()
+
+    plot.stop()
+
+Non-blocking usage::
+
+    plot.start()      # opens window in Qt thread
+    # ... simulation loop calls plot.update() each step
+    plot.stop()       # closes window
+"""
+
+from __future__ import annotations
+
+import sys
+import time
+from collections import deque
+from typing import Any, Callable, Optional
+
+# PyQtGraph and Qt are optional — graceful fallback if not installed
+try:
+    import pyqtgraph as pg
+    from pyqtgraph.Qt import QtWidgets, QtCore
+
+    _HAS_PYQTGRAPH = True
+except ImportError:  # pragma: no cover
+    _HAS_PYQTGRAPH = False
+    pg = None  # type: ignore[assignment]
+    QtWidgets = None  # type: ignore[assignment]
+    QtCore = None  # type: ignore[assignment]
+
+
+class _PlotSeries:
+    """Internal container for one data series."""
+
+    def __init__(
+        self,
+        name: str,
+        source: Callable[[], Any],
+        color: str,
+        max_points: int,
+    ) -> None:
+        self.name = name
+        self.source = source
+        self.color = color
+        self.max_points = max_points
+        self.timestamps: deque[float] = deque(maxlen=max_points)
+        self.values: deque[float] = deque(maxlen=max_points)
+        self.curve: Any = None  # pyqtgraph PlotDataItem
+
+
+class LivePlot:
+    """Real-time multi-trace plotting window powered by PyQtGraph.
+
+    Opens a separate Qt window. The simulation loop must call
+    :meth:`update` periodically to push new data and refresh the display.
+
+    Args:
+        title: Window title.
+        max_points: Maximum number of data points per trace (older points
+            are dropped in a rolling fashion).
+        update_interval_ms: Minimum time between display refreshes in
+            milliseconds. Reduces overhead when :meth:`update` is called
+            faster than needed.
+        y_label: Y-axis label.
+        x_label: X-axis label.
+
+    Example::
+
+        plot = LivePlot(title="Speed vs Time", max_points=300)
+        plot.watch("Speed", lambda: robot.speed, color="#00ff88")
+        plot.start()
+        for _ in range(1000):
+            sim.step()
+            plot.update()
+        plot.stop()
+    """
+
+    def __init__(
+        self,
+        title: str = "BulletLab Live Plot",
+        max_points: int = 500,
+        update_interval_ms: int = 33,
+        y_label: str = "Value",
+        x_label: str = "Time (s)",
+    ) -> None:
+        self._title = title
+        self._max_points = max_points
+        self._update_interval = update_interval_ms / 1000.0
+        self._y_label = y_label
+        self._x_label = x_label
+
+        self._series: list[_PlotSeries] = []
+        self._running = False
+        self._paused = False
+        self._start_time: float = 0.0
+        self._last_refresh: float = 0.0
+
+        # Qt objects (None until start() is called)
+        self._app: Any = None
+        self._window: Any = None
+        self._plot_widget: Any = None
+
+    # ------------------------------------------------------------------
+    # Registration
+    # ------------------------------------------------------------------
+
+    def watch(
+        self,
+        name: str,
+        source: Callable[[], Any],
+        color: str = "#ffffff",
+    ) -> "LivePlot":
+        """Add a data series to the plot.
+
+        Args:
+            name: Series name (shown in legend).
+            source: Callable returning the current value.
+            color: Line color as a hex string (e.g. ``"#00ff88"``).
+
+        Returns:
+            self, for method chaining.
+
+        Example::
+
+            plot.watch("Speed", lambda: robot.speed, color="#00ff88")
+        """
+        series = _PlotSeries(
+            name=name,
+            source=source,
+            color=color,
+            max_points=self._max_points,
+        )
+        self._series.append(series)
+        return self
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(self) -> "LivePlot":
+        """Open the plot window.
+
+        Creates a Qt application if one does not already exist.
+        Non-blocking: the window is opened and control returns immediately.
+        Call :meth:`update` in your simulation loop to refresh the plot.
+
+        Returns:
+            self, for method chaining.
+
+        Raises:
+            ImportError: If PyQtGraph is not installed.
+
+        Example::
+
+            plot.start()
+        """
+        if not _HAS_PYQTGRAPH:
+            print(
+                "[BulletLab] LivePlot: PyQtGraph is not installed. "
+                "Install with: pip install pyqtgraph PyQt5"
+            )
+            return self
+
+        if self._running:
+            return self
+
+        # Ensure Qt application exists
+        self._app = QtWidgets.QApplication.instance()
+        if self._app is None:
+            self._app = QtWidgets.QApplication(sys.argv)
+
+        # Create window
+        self._window = pg.GraphicsLayoutWidget(title=self._title, show=True)
+        self._window.setWindowTitle(self._title)
+        self._window.resize(900, 500)
+
+        self._plot_widget = self._window.addPlot(
+            title=self._title,
+            labels={"left": self._y_label, "bottom": self._x_label},
+        )
+        self._plot_widget.showGrid(x=True, y=True, alpha=0.3)
+        self._plot_widget.addLegend()
+
+        # Create curves for each series
+        for series in self._series:
+            series.curve = self._plot_widget.plot(
+                pen=pg.mkPen(color=series.color, width=2),
+                name=series.name,
+            )
+
+        self._start_time = time.monotonic()
+        self._last_refresh = self._start_time
+        self._running = True
+        return self
+
+    def stop(self) -> None:
+        """Close the plot window and release resources.
+
+        Example::
+
+            plot.stop()
+        """
+        self._running = False
+        if self._window is not None and _HAS_PYQTGRAPH:
+            try:
+                self._window.close()
+            except Exception:
+                pass
+            self._window = None
+
+    # ------------------------------------------------------------------
+    # Update loop
+    # ------------------------------------------------------------------
+
+    def update(self, t: float | None = None) -> None:
+        """Sample all data sources and refresh the plot display.
+
+        Call this once per simulation step. Refresh rate is throttled by
+        ``update_interval_ms`` to avoid excessive Qt overhead.
+
+        Args:
+            t: Timestamp override. If ``None``, uses wall-clock elapsed time
+                since :meth:`start` was called.
+
+        Example::
+
+            for _ in range(5000):
+                sim.step()
+                plot.update()
+        """
+        if not self._running or self._paused:
+            return
+
+        timestamp = t if t is not None else (time.monotonic() - self._start_time)
+
+        # Sample all series
+        for series in self._series:
+            try:
+                val = float(series.source())
+            except Exception:
+                val = float("nan")
+            series.timestamps.append(timestamp)
+            series.values.append(val)
+
+        # Throttle display updates
+        now = time.monotonic()
+        if now - self._last_refresh < self._update_interval:
+            return
+
+        self._last_refresh = now
+        self._refresh_display()
+
+    def _refresh_display(self) -> None:
+        """Update the Qt plot curves from buffered data."""
+        if not _HAS_PYQTGRAPH or self._plot_widget is None:
+            return
+
+        import numpy as np
+
+        for series in self._series:
+            if series.curve is not None and len(series.timestamps) > 0:
+                x = np.array(list(series.timestamps))
+                y = np.array(list(series.values))
+                series.curve.setData(x=x, y=y)
+
+        # Process Qt events
+        try:
+            self._app.processEvents()
+        except Exception:
+            pass
+
+    # ------------------------------------------------------------------
+    # Controls
+    # ------------------------------------------------------------------
+
+    def pause(self) -> None:
+        """Pause data sampling and display updates.
+
+        Example::
+
+            plot.pause()
+        """
+        self._paused = True
+
+    def resume(self) -> None:
+        """Resume a paused plot.
+
+        Example::
+
+            plot.resume()
+        """
+        self._paused = False
+
+    def clear(self) -> None:
+        """Clear all data buffers (but keep series registrations).
+
+        Example::
+
+            plot.clear()
+        """
+        for series in self._series:
+            series.timestamps.clear()
+            series.values.clear()
+
+    def export(self, filepath: str) -> None:
+        """Export the current plot as an image.
+
+        Args:
+            filepath: Output file path. Supported formats: PNG, JPG (via Qt).
+
+        Example::
+
+            plot.export("speed_plot.png")
+        """
+        if not _HAS_PYQTGRAPH or self._window is None:
+            print("[BulletLab] LivePlot: Cannot export — plot window not open.")
+            return
+        try:
+            exporter = pg.exporters.ImageExporter(self._plot_widget)
+            exporter.export(filepath)
+        except Exception as exc:
+            print(f"[BulletLab] LivePlot export failed: {exc}")
+
+    # ------------------------------------------------------------------
+    # State
+    # ------------------------------------------------------------------
+
+    @property
+    def is_running(self) -> bool:
+        """``True`` if the plot window is open."""
+        return self._running
+
+    @property
+    def is_paused(self) -> bool:
+        """``True`` if sampling is paused."""
+        return self._paused
+
+    @property
+    def series_names(self) -> list[str]:
+        """Names of all registered data series."""
+        return [s.name for s in self._series]
+
+    def __repr__(self) -> str:
+        status = "running" if self._running else "stopped"
+        return f"LivePlot({self._title!r}, {status}, series={self.series_names})"

bulletlab/robot/__init__.py

@@ -0,0 +1,12 @@
+"""
+BulletLab robot subpackage.
+
+Provides the Robot, Joint, and Link classes for interacting with simulated
+robots as structured Python objects.
+"""
+
+from bulletlab.robot.robot import Robot
+from bulletlab.robot.joint import Joint, JointType
+from bulletlab.robot.link import Link
+
+__all__ = ["Robot", "Joint", "JointType", "Link"]