robo-goggles 0.1.0__py3-none-any.whl

goggles/history/utils.py

@@ -0,0 +1,191 @@
+"""Utility functions for history slicing and inspection."""
+
+from __future__ import annotations
+
+from typing import Optional, Sequence, Tuple, Union
+
+import jax
+import jaxlib.xla_client as xc
+
+Device = xc.Device
+
+from .types import History
+
+
+def slice_history(
+    history: History,
+    start: int,
+    length: int,
+    fields: Optional[Union[Sequence[str], str]] = None,
+) -> History:
+    """Return a temporal slice [start : start+length] for selected fields.
+
+    Args:
+        history (History): Mapping field -> array of shape (B, T, ...).
+        start (int): Starting timestep (0-based).
+        length (int): Number of timesteps to include (> 0).
+        fields (Optional[Sequence[str] | str]): One or more field names to slice.
+            If a single string is provided, only that field is sliced.
+            If a list or tuple is provided, all listed fields are sliced.
+            If None, all fields in `history` are sliced.
+
+    Returns:
+        History: Mapping of sliced arrays with shape (B, length, ...).
+
+    Raises:
+        ValueError: If `length` <= 0, `start` out of bounds, or slice exceeds T.
+        KeyError: If `fields` is not present in `history`.
+        TypeError: If `history` is empty or contains tensors with rank < 2.
+
+    """
+    # Validate length and history
+    if length <= 0:
+        raise ValueError("length must be > 0")
+    if not history:
+        raise TypeError("history must be a non-empty mapping")
+
+    # Validate reference array and slice bounds
+    any_arr = next(iter(history.values()))
+    if any_arr.ndim < 2:
+        raise TypeError("history arrays must have rank >= 2 (B, T, ...)")
+    T = any_arr.shape[1]
+    if start < 0 or start + length > T:
+        raise ValueError(f"Invalid slice [{start}:{start+length}] for T={T}")
+
+    # Normalize and validate `fields`
+    if fields is None:
+        keys = list(history.keys())
+    elif isinstance(fields, str):
+        keys = [fields]
+    elif isinstance(fields, (list, tuple)):
+        if not fields:
+            raise ValueError("fields list is empty.")
+        if not all(isinstance(f, str) for f in fields):
+            raise TypeError("All field names must be strings.")
+        keys = list(fields)
+    else:
+        raise TypeError("fields must be a string, list/tuple of strings, or None")
+
+    # Check that all requested fields exist
+    missing = set(keys) - set(history)
+    if missing:
+        raise KeyError(f"Unknown fields: {missing}")
+
+    # Validate ranks for selected fields
+    for k in keys:
+        if history[k].ndim < 2:
+            raise TypeError(f"Field {k!r} must have rank >= 2 (B, T, ...)")
+
+    return {k: history[k][:, start : start + length, ...] for k in keys}
+
+
+def peek_last(history: History, k: int = 1) -> History:
+    """Return the last `k` timesteps for all fields.
+
+    Args:
+        history (History): Mapping field -> array of shape (B, T, *payload).
+        k (int): Number of trailing timesteps to select (1 ≤ k ≤ T).
+
+    Returns:
+        History: Mapping field -> sliced array of shape (B, k, *payload).
+
+    Raises:
+        ValueError: If `k` < 1 or `k` > T for any field.
+        TypeError: If `history` is empty or contains tensors with rank < 2.
+
+    """
+    if not history:
+        raise TypeError("history must be a non-empty mapping")
+
+    any_arr = next(iter(history.values()))
+    if any_arr.ndim < 2:
+        raise TypeError("history arrays must have rank >= 2 (B, T, ...)")
+    T = any_arr.shape[1]
+
+    if k < 1 or k > T:
+        raise ValueError(f"k must be in [1, T]; got k={k}, T={T}")
+
+    # Use negative slicing for clarity and to keep JAX-friendly semantics.
+    return {k_name: v[:, -k:, ...] for k_name, v in history.items()}
+
+
+def to_device(
+    history: History,
+    devices: Optional[Sequence[Device]] = None,  # type: ignore[type-arg]
+    keys: Optional[Tuple[str, ...]] = None,
+) -> History:
+    """Move selected history arrays to one or more JAX devices.
+
+    This function moves JAX arrays contained in a dictionary (or subset of it)
+    to a target device or set of devices. If multiple devices are provided,
+    arrays are distributed in a simple round-robin fashion across them.
+
+    Non-array values (e.g., metadata, scalars, strings) are left unchanged.
+
+    Args:
+        history (History): Mapping field to array (or PyTree of arrays).
+        devices (Optional[Sequence[Device]]): Target devices. Defaults to first device.
+        keys (Optional[tuple[str, ...]]): Subset of fields to move. If None, move all.
+
+    Returns:
+        History: Copy of the history with selected arrays placed on the target device(s).
+
+    """
+    devices = devices or jax.devices()
+
+    # Select subset of keys if specified
+    subset = history if keys is None else {k: history[k] for k in keys if k in history}
+
+    moved = {}
+    for i, (k, v) in enumerate(subset.items()):
+        device = devices[i % len(devices)]  # Round-robin device selection
+
+        # Recursively move PyTree leaves to the device
+        moved[k] = jax.tree_util.tree_map(
+            lambda x: jax.device_put(x, device) if isinstance(x, jax.Array) else x,
+            v,
+        )
+
+    # If all keys were moved, just return the moved version
+    if keys is None:
+        return moved
+
+    # Otherwise, merge moved subset back into the original dict
+    return {**history, **moved}
+
+
+def to_host(
+    history: History,
+    keys: Optional[Tuple[str, ...]] = None,
+) -> History:
+    """Copy selected history arrays from device to host memory.
+
+    Recursively retrieves device arrays from JAX devices and copies them
+    into host (NumPy) memory. Non-array values are left unchanged.
+
+    Args:
+        history (History): Mapping field to array (or PyTree of arrays).
+        keys (Optional[tuple[str, ...]]): Subset of fields to copy. If None, all.
+
+    Returns:
+        History: Copy of the history with arrays stored in host (NumPy) memory.
+
+    Example:
+        >>> host_history = to_host(device_history)
+        >>> type(host_history["loss"])
+        <class 'numpy.ndarray'>
+
+    """
+    subset = history if keys is None else {k: history[k] for k in keys if k in history}
+
+    moved = {}
+    for k, v in subset.items():
+        # Recursively copy all arrays from device to host
+        moved[k] = jax.tree_util.tree_map(
+            lambda x: jax.device_get(x) if isinstance(x, jax.Array) else x,
+            v,
+        )
+
+    if keys is None:
+        return moved
+    return {**history, **moved}

goggles/media.py

@@ -0,0 +1,284 @@
+"""Media utilities for saving images and videos from numpy arrays."""
+
+from typing import Literal
+import numpy as np
+import imageio
+from pathlib import Path
+import matplotlib
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+
+def _to_uint8(arr: np.ndarray) -> np.ndarray:
+    if arr.dtype == np.uint8:
+        return arr
+    a = arr.astype(np.float32)
+    vmin = float(np.min(a))
+    vmax = float(np.max(a))
+    if np.isclose(vmax, vmin):
+        return np.zeros_like(a, dtype=np.uint8)
+    if vmin >= 0.0 and vmax <= 1.0:
+        a = a * 255.0
+    else:
+        a = (a - vmin) / (vmax - vmin) * 255.0
+    return np.clip(a, 0, 255).astype(np.uint8)
+
+
+def _normalize_frames(frames: np.ndarray):
+    arr = np.asarray(frames)
+    if arr.ndim < 3:
+        raise ValueError("Expected shape (T, *image_shape[, C]).")
+    # channels
+    if arr.ndim == 3:  # (T,H,W)
+        C = 1
+    else:  # (T,H,W,C)
+        C = arr.shape[-1]
+        if C not in (1, 3):
+            raise ValueError(f"Last dimension must be 1 or 3 channels, got {C}.")
+    if arr.ndim >= 4 and C == 1:
+        arr = arr[..., 0]  # -> (T,H,W)
+    arr_u8 = _to_uint8(arr)
+    mode = "L" if (arr_u8.ndim == 3 or C == 1) else "RGB"
+    return arr_u8, mode
+
+
+def _ensure_even_hw(u8: np.ndarray) -> np.ndarray:
+    if u8.ndim == 3:  # (T,H,W)
+        T, H, W = u8.shape
+        pad_h, pad_w = H % 2, W % 2
+        if pad_h or pad_w:
+            out = np.zeros((T, H + pad_h, W + pad_w), dtype=np.uint8)
+            out[:, :H, :W] = u8
+            return out
+        return u8
+    else:  # (T,H,W,3)
+        T, H, W, C = u8.shape
+        pad_h, pad_w = H % 2, W % 2
+        if pad_h or pad_w:
+            out = np.zeros((T, H + pad_h, W + pad_w, C), dtype=np.uint8)
+            out[:, :H, :W, :] = u8
+            return out
+        return u8
+
+
+def save_numpy_gif(
+    frames: np.ndarray, out_path: str, fps: int = 10, loop: int = 0
+) -> None:
+    """Save a NumPy clip to GIF using imageio.
+
+    Args:
+        frames (np.ndarray): Input clip as a NumPy array of shape
+            (T, H, W) or (T, H, W, C) where C is 1 or 3.
+        out_path (str): Output file path.
+        fps (int): Frames per second.
+        loop (int): Number of times the GIF should loop (0 = infinite).
+
+    """
+    arr_u8, _ = _normalize_frames(frames)
+    imgs = [arr_u8[i] for i in range(arr_u8.shape[0])]
+    imageio.mimsave(
+        out_path,
+        imgs,
+        format="GIF",
+        duration=1.0 / float(fps),
+        loop=loop,
+        palettesize=256,
+        subrectangles=True,
+    )
+
+
+def save_numpy_mp4(
+    frames: np.ndarray,
+    out_path: Path,
+    fps: int = 30,
+    codec: str = "libx264",
+    pix_fmt: str = "yuv420p",
+    bitrate: str | None = None,
+    crf: int | None = 18,
+    convert_gray_to_rgb: bool = True,
+    preset: str | None = "medium",
+) -> None:
+    """Save a NumPy clip to MP4 using imageio-ffmpeg.
+
+    Args:
+        frames (np.ndarray): Input clip as a NumPy array of shape
+            (T, H, W) or (T, H, W, C) where C is 1 or 3.
+        out_path (str): Output file path.
+        fps (int): Frames per second.
+        codec (str): Video codec to use.
+        pix_fmt (str): Pixel format for ffmpeg.
+        bitrate (str | None): Bitrate string for ffmpeg (e.g. "4M").
+        crf (int | None): Constant Rate Factor for quality control.
+        convert_gray_to_rgb (bool): Whether to convert grayscale to RGB.
+        preset (str | None): ffmpeg preset for speed/quality tradeoff.
+
+    """
+    arr_u8, mode = _normalize_frames(frames)
+
+    # Prefer RGB for broad player compatibility
+    if mode == "L" and convert_gray_to_rgb:
+        arr_u8 = np.stack([arr_u8] * 3, axis=-1)  # (T,H,W) -> (T,H,W,3)
+
+    # Ensure even dims for yuv420p
+    arr_u8 = _ensure_even_hw(arr_u8)
+
+    # Build writer kwargs
+    writer_kwargs = {
+        "fps": fps,
+        "codec": codec,
+        "macro_block_size": None,
+        "format": "FFMPEG",
+    }
+    if bitrate is not None:
+        writer_kwargs["bitrate"] = bitrate
+
+    ffmpeg_params = []
+    if crf is not None:
+        ffmpeg_params += ["-crf", str(crf)]
+    if preset is not None:
+        ffmpeg_params += ["-preset", preset]
+    if pix_fmt is not None:
+        ffmpeg_params += ["-pix_fmt", pix_fmt]
+    if ffmpeg_params:
+        writer_kwargs["ffmpeg_params"] = ffmpeg_params
+
+    with imageio.get_writer(out_path, **writer_kwargs) as writer:
+        if arr_u8.ndim == 3:
+            for i in range(arr_u8.shape[0]):
+                writer.append_data(arr_u8[i])
+        else:
+            for i in range(arr_u8.shape[0]):
+                writer.append_data(arr_u8[i])
+
+
+def save_numpy_image(image: np.ndarray, out_path: str, format: str) -> None:
+    """Save a NumPy image to file using imageio.
+
+    Args:
+        image (np.ndarray): Input image as a NumPy array of shape
+            (H, W) or (H, W, C) where C is 1 or 3.
+        out_path (str): Output file path.
+        format (str): Image format (e.g., 'png', 'jpg', 'jpeg').
+
+    """
+    arr_u8, _ = _normalize_frames(image[np.newaxis, ...])
+    imageio.imwrite(out_path, arr_u8[0], format=format)
+
+
+def save_numpy_vector_field_visualization(
+    vector_field: np.ndarray,
+    dir: Path,
+    name: str,
+    mode: Literal["vorticity", "magnitude"] = "magnitude",
+    arrow_stride: int = 8,
+    dpi: int = 300,
+    add_colorbar: bool = True,
+) -> None:
+    """Save a 2D vector field visualization as a PNG image.
+
+    Args:
+        vector_field (np.ndarray): Input vector field of shape (H, W, 2).
+        dir (Path): Output directory path.
+        name (str): Base name for the output PNG file (without extension).
+        mode (Literal["vorticity", "magnitude"]): Visualization mode.
+        arrow_stride (int): Stride for downsampling arrows (every Nth point).
+        dpi (int): Resolution of the output image.
+        add_colorbar (bool): Whether to include a colorbar.
+
+    """
+    # Store original backend to restore later
+    original_backend = matplotlib.get_backend()
+    matplotlib.use("Agg")
+
+    try:
+        H, W, _ = vector_field.shape
+
+        # Create figure that matches pixel aspect; keep margins zero by default
+        fig, ax = plt.subplots(figsize=(W / 50, H / 50), dpi=dpi)
+        ax.set_aspect("equal")
+
+        # Compute scalar field and arrow color based on the selected mode
+        if mode == "magnitude":
+            scalar_field = np.linalg.norm(vector_field, axis=-1)
+            cmap = plt.cm.viridis
+            arrow_color = "white"
+        elif mode == "vorticity":
+            # Compute vorticity (curl) of the vector field: dVx/dy - dVy/dx
+            dy = np.gradient(vector_field[..., 0], axis=0)
+            dx = np.gradient(vector_field[..., 1], axis=1)
+            scalar_field = dx - dy
+            cmap = plt.cm.RdBu_r
+            arrow_color = "black"
+        else:
+            raise ValueError("mode must be 'magnitude' or 'vorticity'")
+
+        # Display scalar field as background
+        im = ax.imshow(
+            scalar_field,
+            cmap=cmap,
+            origin="lower",
+            extent=[0, W, 0, H],
+            interpolation="bilinear",
+        )
+
+        # Arrow grid
+        y_coords, x_coords = np.mgrid[0:H:arrow_stride, 0:W:arrow_stride]
+        u_sampled = vector_field[::arrow_stride, ::arrow_stride, 0]
+        v_sampled = vector_field[::arrow_stride, ::arrow_stride, 1]
+
+        # Plot arrows
+        ax.quiver(
+            x_coords,
+            y_coords,
+            u_sampled,
+            v_sampled,
+            color=arrow_color,
+            alpha=0.9,
+            scale_units="xy",
+            scale=1,
+            width=0.002,
+            headwidth=4,
+            headlength=5,
+            headaxislength=4.5,
+            linewidth=0.5,
+            edgecolor="none",
+        )
+
+        # Remove axes and ticks
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.set_frame_on(False)
+
+        # (1) Colorbar exactly same height as the axes
+        if add_colorbar:
+            divider = make_axes_locatable(ax)
+            # size can be tweaked; "3%" is a nice thin bar, pad is the gap
+            cax = divider.append_axes("right", size="3%", pad=0.02)
+            cb = fig.colorbar(im, cax=cax)
+            cb.ax.tick_params(length=2)
+
+            # Leave a tiny margin so ticks/labels aren't clipped
+            fig.subplots_adjust(left=0.0, right=0.98, bottom=0.0, top=1.0)
+            bbox_setting = "tight"
+            pad_setting = 0.02
+        else:
+            # (2) No colorbar: strip all outer boundaries/margins
+            fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
+            bbox_setting = "tight"
+            pad_setting = 0.0
+
+        # Save
+        dir.mkdir(parents=True, exist_ok=True)
+        plt.savefig(
+            str(dir / f"{name}.png"),
+            dpi=dpi,
+            bbox_inches=bbox_setting,
+            pad_inches=pad_setting,
+            facecolor="white",
+            edgecolor="none",
+        )
+        plt.close(fig)
+
+    finally:
+        # Restore original matplotlib backend
+        matplotlib.use(original_backend)

goggles/shutdown.py

@@ -0,0 +1,70 @@
+"""Simple util for graceful shutdowns in Python applications."""
+
+import signal
+from typing import Optional
+
+
+class GracefulShutdown:
+    """A context manager for graceful shutdowns.
+
+    Example:
+    >>> with GracefulShutdown(exit_message="Shutting down gracefully...") as gs:
+    ...     while not gs.stop:
+    ...         # Main application logic here, runs until interrupted
+    ...         # by SIGINT or SIGTERM.
+    ...         pass
+    ...     print("Cleanup and exit.")
+
+    """
+
+    stop = False
+
+    def __init__(
+        self,
+        exit_message: Optional[str] = None,
+    ):
+        """Initialize the GracefulShutdown context manager.
+
+        Args:
+            exit_message (str): The message to log upon shutdown.
+
+        """
+        from . import get_logger
+
+        self.logger = get_logger("goggles.shutdown")
+        self.exit_message = exit_message
+        # placeholders for original handlers
+        self._orig_sigint = None
+        self._orig_sigterm = None
+
+    def __enter__(self):
+        """Register the signal handlers."""
+        # save existing handlers
+        self._orig_sigint = signal.getsignal(signal.SIGINT)
+        self._orig_sigterm = signal.getsignal(signal.SIGTERM)
+
+        def handle_signal(signum, frame):
+            self.stop = True
+            if self.exit_message:
+                self.logger.info(self.exit_message)
+
+        # register for both SIGINT and SIGTERM
+        signal.signal(signal.SIGINT, handle_signal)
+        signal.signal(signal.SIGTERM, handle_signal)
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Unregister the signal handlers, restoring originals.
+
+        Args:
+            exc_type: Exception type if any.
+            exc_value: Exception value if any.
+            traceback: Traceback if any.
+
+        """
+        # restore original handlers
+        if self._orig_sigint is not None:
+            signal.signal(signal.SIGINT, self._orig_sigint)
+        if self._orig_sigterm is not None:
+            signal.signal(signal.SIGTERM, self._orig_sigterm)

goggles/types.py

@@ -0,0 +1,79 @@
+"""Types used in Goggles."""
+
+import numpy as np
+from typing import Dict, Literal, Any, Optional
+from dataclasses import dataclass
+from typing import TypeAlias
+
+Kind = Literal["log", "metric", "image", "video", "artifact"]
+
+Metrics = Dict[str, float | int]
+Image: TypeAlias = np.ndarray
+Video: TypeAlias = np.ndarray
+Vector: TypeAlias = np.ndarray
+VectorField: TypeAlias = np.ndarray
+
+
+@dataclass(frozen=True)
+class Event:
+    """Structured event routed through the EventBus.
+
+    Args:
+        kind (Kind): Kind of event ("log", "metric", "image", "artifact").
+        scope (str): Scope of the event ("global" or "run").
+        payload (Any): Event payload.
+        filepath (str): File path of the caller emitting the event.
+        lineno (int): Line number of the caller emitting the event.
+        level (Optional[int]): Optional log level for "log" events.
+        step (Optional[int]): Optional global step index.
+        time (Optional[float]): Optional global timestamp.
+        extra (Optional[dict[str, Any]]): Optional extra metadata.
+
+    """
+
+    kind: Kind
+    scope: str
+    payload: Any
+    filepath: str
+    lineno: int
+    level: Optional[int] = None
+    step: Optional[int] = None
+    time: Optional[float] = None
+    extra: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert Event to dictionary."""
+        result = {
+            "kind": self.kind,
+            "scope": self.scope,
+            "payload": self.payload,
+            "filepath": self.filepath,
+            "lineno": self.lineno,
+        }
+
+        # Only include optional fields if they are not None
+        if self.level is not None:
+            result["level"] = self.level
+        if self.step is not None:
+            result["step"] = self.step
+        if self.time is not None:
+            result["time"] = self.time
+        if self.extra is not None:
+            result["extra"] = self.extra
+
+        return result
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Event":
+        """Create Event from dictionary."""
+        return cls(
+            kind=data["kind"],
+            scope=data["scope"],
+            payload=data["payload"],
+            filepath=data["filepath"],
+            lineno=data["lineno"],
+            level=data.get("level"),
+            step=data.get("step"),
+            time=data.get("time"),
+            extra=data.get("extra"),
+        )