Add side_by_side() for horizontal animation concatenation

streamjoy/__init__.py

@@ -1,15 +1,15 @@
 import logging
 
 from ._utils import update_logger
-from .core import connect, stream
+from .core import connect, side_by_side, stream
 from .models import ImageText, Paused
 from .renderers import (
     default_holoviews_renderer,
     default_pandas_renderer,
     default_xarray_renderer,
 )
 from .settings import config, file_handlers, obj_handlers
-from .streams import GifStream, HtmlStream, Mp4Stream
+from .streams import GifStream, HtmlStream, Mp4Stream, SideBySideStreams
 from .wrappers import wrap_holoviews, wrap_matplotlib
 
 __version__ = "0.0.10"
@@ -20,13 +20,15 @@
     "ImageText",
     "Mp4Stream",
     "Paused",
+    "SideBySideStreams",
     "config",
     "connect",
     "default_holoviews_renderer",
     "default_pandas_renderer",
     "default_xarray_renderer",
     "file_handlers",
     "obj_handlers",
+    "side_by_side",
     "stream",
     "wrap_holoviews",
     "wrap_matplotlib",

streamjoy/core.py

@@ -1,13 +1,21 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from io import BytesIO
 from pathlib import Path
-from typing import Any, Callable, Literal
+from typing import Any, Literal
 
 from . import streams
 from .serializers import serialize_appropriately
 from .settings import extension_handlers
-from .streams import AnyStream, ConnectedStreams, GifStream, HtmlStream, Mp4Stream
+from .streams import (
+    AnyStream,
+    ConnectedStreams,
+    GifStream,
+    HtmlStream,
+    Mp4Stream,
+    SideBySideStreams,
+)
 
 
 def stream(
@@ -79,3 +87,29 @@ def connect(
     if uri:
         return stream.write(uri=uri)
     return stream
+
+
+def side_by_side(
+    streams: list[AnyStream | GifStream | Mp4Stream | HtmlStream],
+    uri: str | Path | BytesIO | None = None,
+) -> SideBySideStreams | Path | BytesIO:
+    """
+    Render heterogeneous streams side by side (horizontally concatenated).
+
+    Frames from each stream are combined horizontally to create a single
+    animation with multiple streams playing simultaneously. If streams have
+    different numbers of frames, the shorter streams will repeat their last
+    frame to match the longest stream.
+
+    Args:
+        streams: The streams to render side by side.
+        uri: The destination to write the side-by-side streams to.
+            If None, the side-by-side streams object is returned.
+
+    Returns:
+        The side-by-side streams if uri is None, otherwise the uri (Path or BytesIO).
+    """
+    stream = SideBySideStreams(streams=streams)
+    if uri:
+        return stream.write(uri=uri)
+    return stream

streamjoy/streams.py

@@ -3,14 +3,14 @@
 import base64
 import gc
 from abc import abstractmethod
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from contextlib import contextmanager
 from functools import partial
 from io import BytesIO
 from itertools import zip_longest
 from pathlib import Path
 from textwrap import indent
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 import dask.delayed
 import imageio.v3 as iio
@@ -69,11 +69,7 @@ class MediaStream(param.Parameterized):
     Expand Source code to see all the parameters and descriptions.
     """
 
-    resources = param.List(
-        default=None,
-        doc="The resources to render.",
-        precedence=0
-    )
+    resources = param.List(default=None, doc="The resources to render.", precedence=0)
 
     renderer = param.Callable(
         doc="The renderer to use for the resources.",
@@ -970,7 +966,7 @@ def _optimize_gif(self, uri: Path) -> None:
         green = config["logging_success_color"]
         reset = config["logging_reset_color"]
         self.logger.success(
-            f"Optimized stream with pygifsicle at " f"{green}{uri.absolute()}{reset}."
+            f"Optimized stream with pygifsicle at {green}{uri.absolute()}{reset}."
         )
 
     def write(
@@ -1371,3 +1367,204 @@ def __iter__(self) -> Iterable[Path]:
         Iterate over the streams.
         """
         return self.streams
+
+
+class SideBySideStreams(param.Parameterized):
+    """
+    Multiple streams rendered side by side (horizontally concatenated).
+    """
+
+    streams = param.List(
+        default=[],
+        item_type=MediaStream,
+        doc="The streams to render side by side.",
+    )
+
+    def __init__(self, streams: list[MediaStream] | None = None, **params) -> None:
+        params["streams"] = streams or params.get("streams")
+        self.logger = _utils.update_logger()
+        super().__init__(**params)
+
+    def write(
+        self,
+        uri: str | Path | BytesIO | None = None,
+        extension: str | None = None,
+        **kwargs: dict[str, Any],
+    ) -> Path | BytesIO:
+        """
+        Write the side-by-side streams to a file or in memory.
+
+        Args:
+            uri: The file path or BytesIO object to write to.
+            extension: The file extension to use.
+            **kwargs: Additional keyword arguments to pass.
+
+        Returns:
+            The file path or BytesIO object.
+        """
+        if uri is None:
+            uri = BytesIO()
+        elif isinstance(uri, str):
+            uri = Path(uri)
+
+        streams = [
+            (
+                stream.materialize(uri, extension)
+                if isinstance(stream, AnyStream)
+                else stream
+            )
+            for stream in self.streams
+        ]
+
+        client = _utils.get_distributed_client(
+            client=kwargs.get("client"),
+            processes=kwargs.get("processes"),
+            threads_per_worker=kwargs.get("threads_per_worker"),
+        )
+
+        stream_0 = streams[0]
+        extension = stream_0._extension
+
+        # Render all stream images
+        all_images = []
+        for stream in streams:
+            stream.client = client
+            resources, renderer_iterables = _utils.subset_resources_renderer_iterables(
+                stream.resources, stream.renderer_iterables, stream.max_frames
+            )
+            images = stream._render_images(
+                resources,
+                stream.renderer,
+                renderer_iterables,
+                stream.renderer_kwargs,
+            )
+            all_images.append(images)
+
+        # Get the maximum number of frames across all streams
+        max_frames = max(len(images) for images in all_images)
+
+        # Combine frames side by side
+        combined_images = []
+        for frame_idx in range(max_frames):
+            frame_arrays = []
+            max_height = 0
+
+            # Collect frames from each stream for this index
+            for stream_idx, images in enumerate(all_images):
+                if frame_idx < len(images):
+                    # Get the frame
+                    frame = _utils.get_result(images[frame_idx])
+                else:
+                    # Use last frame if this stream has fewer frames
+                    frame = _utils.get_result(images[-1])
+
+                # Handle Paused frames
+                if isinstance(frame, Paused):
+                    frame = frame.output
+
+                # Convert to numpy array if it's a PIL Image
+                if isinstance(frame, Image.Image):
+                    frame = np.array(frame)
+
+                frame_arrays.append(frame)
+                max_height = max(max_height, frame.shape[0])
+
+            # Pad frames to same height if needed
+            padded_frames = []
+            for frame in frame_arrays:
+                if frame.shape[0] < max_height:
+                    # Pad vertically with black
+                    pad_height = max_height - frame.shape[0]
+                    if len(frame.shape) == 3:
+                        # RGB/RGBA image
+                        padding = np.zeros(
+                            (pad_height, frame.shape[1], frame.shape[2]),
+                            dtype=frame.dtype,
+                        )
+                    else:
+                        # Grayscale image
+                        padding = np.zeros(
+                            (pad_height, frame.shape[1]), dtype=frame.dtype
+                        )
+                    frame = np.vstack([frame, padding])
+                padded_frames.append(frame)
+
+            # Concatenate horizontally
+            combined_frame = np.hstack(padded_frames)
+            combined_images.append(combined_frame)
+
+        # Write combined frames to output
+        duration = []
+        if isinstance(stream_0, GifStream):
+            # Calculate duration for all frames
+            duration = [int(1000 / stream_0.fps)] * len(combined_images)
+
+        write_kwargs = stream_0.write_kwargs.copy()
+
+        with iio.imopen(uri, "w", extension=extension) as buf:
+            if isinstance(stream_0, GifStream):
+                write_kwargs["duration"] = duration
+                for frame in combined_images:
+                    buf.write(frame, **write_kwargs)
+            elif isinstance(stream_0, Mp4Stream):
+                # Initialize video stream for MP4
+                init_kwargs = _utils.pop_kwargs(buf.init_video_stream, write_kwargs)
+                init_kwargs["fps"] = stream_0.fps
+                buf.init_video_stream(stream_0.codec, **init_kwargs)
+
+                if "crf" in write_kwargs:
+                    buf._video_stream.options = {"crf": str(write_kwargs.pop("crf"))}
+
+                for frame in combined_images:
+                    # Ensure dimensions are even for MP4
+                    if frame.shape[0] % 2:
+                        frame = frame[:-1, :, :]
+                    if frame.shape[1] % 2:
+                        frame = frame[:, :-1, :]
+                    frame = frame[:, :, :3]  # remove alpha channel if present
+                    buf.write_frame(frame, **write_kwargs)
+            else:
+                # For other stream types
+                for i, frame in enumerate(combined_images):
+                    write_kwargs["init"] = i == 0
+                    buf.write(frame, **write_kwargs)
+
+        # Optimize GIF if needed
+        if isinstance(stream_0, GifStream) and stream_0.optimize:
+            stream_0._optimize_gif(uri)
+
+        green = config["logging_success_color"]
+        reset = config["logging_reset_color"]
+        if isinstance(uri, Path):
+            self.logger.success(f"Saved stream to {green}{uri.absolute()}{reset}.")
+        else:
+            self.logger.success(f"Saved stream to {green}memory{reset}.")
+        stream_0._display_in_notebook(uri, display=stream_0.display)
+        return uri
+
+    def __repr__(self) -> str:
+        repr_str = (
+            f"<{self.__class__.__name__}>\n"
+            f"---\n"
+            f"Streams (side by side):\n"
+            f"{indent(repr(self.streams[0]), ' ' * 2)}\n"
+            f"  ...\n"
+            f"{indent(repr(self.streams[-1]), ' ' * 2)}\n"
+        )
+        return repr_str
+
+    def __len__(self) -> int:
+        """
+        Return the maximum number of frames across all streams.
+
+        Note: This differs from ConnectedStreams which returns the sum of frames.
+        For side-by-side streams, all streams play simultaneously, so the length
+        is the maximum frame count, not the sum.
+        """
+        return max(len(stream) for stream in self.streams) if self.streams else 0
+
+    def __iter__(self) -> Iterable[MediaStream]:
+        """
+        Iterate over the streams.
+        """
+        return self.streams