parallel-matplotlib-animation 0.1.0__py3-none-any.whl

parallel_animate/__init__.py

@@ -0,0 +1 @@
+from .animator import Animator

parallel_animate/animator.py

@@ -0,0 +1,387 @@
+import matplotlib
+
+matplotlib.use("Agg")
+
+import tempfile
+import logging
+import time
+import multiprocessing as mp
+from pathlib import Path
+from abc import ABC, abstractmethod
+from typing import Any
+
+import matplotlib.pyplot as plt
+import av
+from PIL import Image
+from tqdm import tqdm, trange
+
+
+class Animator(ABC):
+    """
+    Base class for creating matplotlib animations with efficient parallel rendering.
+    """
+
+    def __init__(self):
+        """Initialize the animator."""
+        self.logger = logging.getLogger(__name__)
+
+    @abstractmethod
+    def setup(self):
+        """
+        Set up the figure, axes, and any artists. Store whatever you need as instance
+        attributes (self.ax, self.line, etc.); you will need them in update().
+
+        This method MUST not accept any argument other than self (if you need to pass
+        data to setup, store it as instance attributes in __init__).
+
+        This method MUST return the plt.Figure object, and that alone.
+
+        This method is called once before the animation loop (once per worker if
+        parallelized), unless reuse_figure_object is False (basically never).
+        """
+        pass
+
+    @abstractmethod
+    def update(self, frame_idx: int, params: Any):
+        """
+        Update the plot for the given frame.
+
+        This is called for each frame in the animation. Modify any figure, axes, and
+        artists that you created in setup() and stored as instance attributes.
+
+        Args:
+            frame_idx (int): Current frame index (0 to num_frames-1)
+            params (Any): Parameters for this frame from param_by_frame[frame_idx]. Can
+                be any object you want as long as it's pickleable (required to
+                distribute jobs to distributed workers).
+        """
+        pass
+
+    def make_video(
+        self,
+        output_file: Path | str,
+        param_by_frame: list[Any],
+        fps: int,
+        num_workers: int = -1,
+        disable_progress_bar: bool | None = None,
+        plotting_log_interval: int | None = None,
+        saving_log_interval: int | None = None,
+        savefig_params: dict[str, Any] = {},
+        video_codec: str = "libx264",
+        video_params: dict[str, Any] = {"pix_fmt": "yuv420p"},
+        reuse_figure_object: bool = True,
+    ) -> None:
+        """
+        Render the animation to a video file.
+
+        Args:
+            output_file (Path | str): Path to output video file
+            param_by_frame (list[Any]): List of parameters, one per frame
+            fps (int): Frames per second for output video
+            num_workers (int): Number of parallel workers. 1 for serial processing
+                (in the main thread), -1 for all CPU cores, -2 for all but one CPU core,
+                etc.
+            disable_progress_bar (bool | None): Same behavior as tqdm: if True, disable
+                progress bar. If None, auto-detect based on whether output is a TTY.
+            plotting_log_interval (int | None): Log progress every N frames in each
+                worker (None = no interval logging).
+            saving_log_interval (int | None): Log progress every N frames when merging
+                frames into video (None = no interval logging).
+            savefig_params (dict[str, Any]): Additional keyword arguments to
+                pass to plt.Figure.savefig() when saving frames (default: {}).
+            video_codec (str): Codec to use for video encoding (default: "libx264").
+            video_params (dict[str, Any]): Additional parameters to set on the video
+                stream (default: {"pix_fmt": "yuv420p"}).
+            reuse_figure_object (bool): If False, the figure will be re-created for each
+                frame (i.e. setup() called every frame). There is basically no reason to
+                set this to False. Use only for testing and benchmarking.
+        """
+        # Try to convert param_by_frame to list
+        try:
+            params_list = list(param_by_frame)
+        except Exception as e:
+            self.logger.critical(
+                "param_by_frame must be convertible to a list. Ensure it is a list-like object."
+            )
+            raise e
+
+        num_frames = len(params_list)
+
+        # Determine number of workers
+        if num_workers == -1:
+            num_workers = mp.cpu_count()
+        elif num_workers < -1:
+            num_workers = max(1, mp.cpu_count() + num_workers + 1)
+
+        self.logger.info(f"Rendering {num_frames} frames at {fps} fps")
+        with tempfile.TemporaryDirectory(prefix="animator_frames_") as frames_dir:
+            self.logger.info(f"Using temporary directory: {frames_dir}")
+
+            if num_workers == 1:
+                self.logger.info("Running in serial mode")
+                self._render_serial(
+                    params_list,
+                    frames_dir,
+                    disable_progress_bar,
+                    plotting_log_interval,
+                    savefig_params,
+                    reuse_figure_object,
+                )
+            else:
+                self.logger.info(f"Running in parallel mode with {num_workers} workers")
+                self._render_parallel(
+                    params_list,
+                    frames_dir,
+                    num_workers,
+                    disable_progress_bar,
+                    plotting_log_interval,
+                    savefig_params,
+                    reuse_figure_object,
+                )
+
+            self.logger.info("Creating video with PyAV")
+            _merge_frames_into_video(
+                frames_dir,
+                output_file,
+                fps,
+                video_codec,
+                video_params,
+                disable_progress_bar,
+                self.logger,
+                log_interval=saving_log_interval,
+            )
+
+        self.logger.info(f"Animation complete: {output_file}")
+
+    def _setup_and_check(self) -> plt.Figure:
+        """Call setup() and validate its return type."""
+        fig = self.setup()
+        if not isinstance(fig, plt.Figure):
+            raise TypeError(
+                f"`.setup()` must return a matplotlib Figure object, got {type(fig).__name__}"
+            )
+        return fig
+
+    def _render_serial(
+        self,
+        params_list: list[Any],
+        frames_dir: Path | str,
+        disable_progress_bar: bool | None,
+        log_interval: int | None,
+        savefig_params: dict[str, Any],
+        reuse_figure_object: bool,
+    ) -> None:
+        """Render frames serially."""
+        self.logger.info("Serial rendering")
+
+        if reuse_figure_object:
+            # Setup once and get figure
+            fig = self._setup_and_check()
+
+        # Render all frames with progress bar
+        num_frames = len(params_list)
+        for frame_idx in trange(
+            num_frames, desc="Rendering", disable=disable_progress_bar
+        ):
+            if not reuse_figure_object:
+                fig = self._setup_and_check()
+
+            params = params_list[frame_idx]
+            self.update(frame_idx, params)
+            fig.canvas.draw()
+
+            frame_path = Path(frames_dir) / f"frame_{frame_idx:09d}.png"
+            fig.savefig(frame_path, **savefig_params)
+            if not reuse_figure_object:
+                plt.close(fig)
+
+            # Optional interval logging
+            if log_interval and (frame_idx + 1) % log_interval == 0:
+                self.logger.info(f"Frame {frame_idx + 1}/{num_frames}")
+
+        plt.close(fig)
+
+    def _render_parallel(
+        self,
+        params_list: list[Any],
+        frames_dir: Path | str,
+        num_workers: int,
+        disable_progress_bar: bool | None,
+        log_interval: int | None,
+        savefig_params: dict[str, Any],
+        reuse_figure_object: bool,
+    ) -> None:
+        """Render frames in parallel using dynamic work distribution."""
+        num_frames = len(params_list)
+
+        self.logger.info(f"Using dynamic work distribution with {num_workers} workers")
+
+        # Create queues for task distribution and atomic counter for progress
+        task_queue = mp.Queue()
+        num_frames_completed = mp.Value("i", 0)  # atomic integer counter
+
+        # Start worker processes
+        workers = []
+        for worker_id in range(num_workers):
+            p = mp.Process(
+                target=_worker_process,
+                args=(
+                    self,
+                    worker_id,
+                    task_queue,
+                    num_frames_completed,
+                    frames_dir,
+                    log_interval,
+                    savefig_params,
+                    reuse_figure_object,
+                ),
+            )
+            p.start()
+            workers.append(p)
+
+        # Populate task queue with individual frames (batch_size = 1)
+        for frame_idx in range(num_frames):
+            task_queue.put((frame_idx, params_list[frame_idx]))
+
+        # Send sentinel values to signal workers to exit
+        for _ in range(num_workers):
+            task_queue.put(None)
+
+        # Monitor progress using atomic counter
+        pbar = tqdm(total=num_frames, desc="Rendering", disable=disable_progress_bar)
+        while num_frames_completed.value < num_frames:
+            current_progress = num_frames_completed.value
+            if current_progress > pbar.n:
+                pbar.update(current_progress - pbar.n)
+            time.sleep(0.1)
+        pbar.update(num_frames_completed.value - pbar.n)
+        pbar.close()
+
+        # Wait for all workers to finish
+        for p in workers:
+            p.join()
+
+        self.logger.info("All workers completed")
+
+
+def _worker_process(
+    animator: Animator,
+    worker_id: int,
+    task_queue: mp.Queue,
+    progress_counter,
+    frames_dir: Path | str,
+    log_interval: int | None,
+    savefig_params: dict[str, Any],
+    reuse_figure_object: bool,
+) -> None:
+    """
+    Worker process that renders frames.
+
+    Each worker:
+    1. Calls setup() once to initialize the figure (unless reuse_figure_object is False)
+    2. Repeatedly pulls individual frames from the task queue
+    3. Renders each frame
+    4. Atomically increments the progress counter
+    """
+    # Setup once per worker
+    if reuse_figure_object:
+        fig = animator._setup_and_check()
+
+    # Process frames until we get a sentinel value (None)
+    frames_processed = 0
+    while True:
+        task = task_queue.get()
+        if task is None:
+            break  # sentinel value - exit
+        frame_idx, params = task
+
+        # Render the frame
+        if not reuse_figure_object:
+            fig = animator._setup_and_check()
+        animator.update(frame_idx, params)
+        fig.canvas.draw()
+        frame_path = Path(frames_dir) / f"frame_{frame_idx:09d}.png"
+        fig.savefig(frame_path, **savefig_params)
+        frames_processed += 1
+        if not reuse_figure_object:
+            plt.close(fig)
+
+        # Logging
+        if log_interval and frames_processed % log_interval == 0:
+            animator.logger.info(
+                f"Worker {worker_id}: processed {frames_processed} frames"
+            )
+
+        # Atomically increment progress counter
+        with progress_counter.get_lock():
+            progress_counter.value += 1
+
+    plt.close(fig)
+    animator.logger.info(f"Worker {worker_id}: completed {frames_processed} frames")
+
+
+def _merge_frames_into_video(
+    frames_dir: Path | str,
+    output_file: Path | str,
+    fps: int,
+    video_codec: str,
+    video_params: dict[str, Any],
+    disable_progress_bar: bool | None,
+    logger: logging.Logger,
+    log_interval: int | None,
+) -> None:
+    """Use PyAV to merge frames into video."""
+    # Gather frame files in sorted order
+    frame_files = sorted(Path(frames_dir).glob("frame_*.png"))
+
+    if not frame_files:
+        raise RuntimeError("No frames found in temporary directory")
+
+    # Open first image to determine size and ensure even dimensions
+    with Image.open(frame_files[0]) as first_img:
+        width, height = first_img.size
+    # Make width/height even (required by many codecs)
+    width = (width // 2) * 2
+    height = (height // 2) * 2
+
+    try:
+        Path(output_file).parent.mkdir(parents=True, exist_ok=True)
+        container = av.open(str(output_file), mode="w")
+        stream = container.add_stream(video_codec, rate=fps)
+        stream.width = width
+        stream.height = height
+        for key, value in video_params.items():
+            setattr(stream, key, value)
+
+        # Encode each frame
+        for i, frame_path in tqdm(
+            enumerate(frame_files),
+            total=len(frame_files),
+            desc="Merging frames",
+            disable=disable_progress_bar,
+        ):
+            img = Image.open(frame_path).convert("RGBA")
+            # Ensure image has the right size
+            if img.size != (width, height):
+                img = img.resize((width, height))
+
+            video_frame = av.VideoFrame.from_image(img)
+            for packet in stream.encode(video_frame):
+                container.mux(packet)
+
+            # Optional interval logging
+            if log_interval and (i + 1) % log_interval == 0:
+                logger.info(f"Frame written {i + 1}/{len(frame_files)} to video")
+
+        # Flush encoder
+        for packet in stream.encode(None):
+            container.mux(packet)
+
+        container.close()
+
+        size_mb = Path(output_file).stat().st_size / (1024 * 1024)
+        logger.info(f"Video created: {output_file} ({size_mb:.2f} MB)")
+
+    except Exception as e:
+        logger.critical(f"PyAV failed to create video: {e}")
+        raise e

parallel_animate/examples/multi_panel_animation.py

@@ -0,0 +1,144 @@
+from pathlib import Path
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+from parallel_animate import Animator
+
+
+class MultiPanelAnimation(Animator):
+    def __init__(self):
+        super().__init__()
+
+    def setup(self):
+        # Create figure with 1 row, 5 columns of subplots
+        self.fig, axes = plt.subplots(1, 5, figsize=(20, 4))
+        self.axes = list(axes)
+
+        # Replace 4th subplot with 3D version
+        self.axes[3].remove()
+        self.axes[3] = self.fig.add_subplot(1, 5, 4, projection="3d")
+
+        # Initialize artists for each subplot
+        self.artists = {}
+
+        # Line chart - multiple overlapping waves
+        x = np.linspace(0, 4 * np.pi, 100)
+        (line,) = self.axes[0].plot(x, np.sin(x), "b-", linewidth=2)
+        self.axes[0].set_xlim(0, 4 * np.pi)
+        self.axes[0].set_ylim(-2, 2)
+        self.axes[0].set_title("Line Chart")
+        self.axes[0].grid(True, alpha=0.3)
+        self.artists["line"] = line
+        self.artists["line_x"] = x
+
+        # Bar chart - dancing bars
+        bars = self.axes[1].bar(range(10), np.random.rand(10), color="steelblue")
+        self.axes[1].set_ylim(0, 2)
+        self.axes[1].set_title("Bar Chart")
+        self.axes[1].set_xlabel("Category")
+        self.axes[1].set_ylabel("Value")
+        self.artists["bars"] = bars
+
+        # Imshow - wave interference pattern
+        data = np.random.rand(30, 30)
+        im = self.axes[2].imshow(data, cmap="viridis", vmin=0, vmax=2)
+        self.axes[2].set_title("Wave Interference")
+        self.axes[2].axis("off")
+        self.fig.colorbar(im, ax=self.axes[2], fraction=0.046, pad=0.04)
+        self.artists["imshow"] = im
+
+        # 3D line - rotating helix
+        t = np.linspace(0, 4 * np.pi, 100)
+        (line3d,) = self.axes[3].plot(np.cos(t), np.sin(t), t, "b-", linewidth=2)
+        self.axes[3].set_xlim(-2, 2)
+        self.axes[3].set_ylim(-2, 2)
+        self.axes[3].set_zlim(0, 4 * np.pi)
+        self.axes[3].set_title("3D Helix")
+        self.axes[3].set_xlabel("X")
+        self.axes[3].set_ylabel("Y")
+        self.axes[3].set_zlabel("Z")
+        self.artists["line3d"] = line3d
+        self.artists["line3d_t"] = t
+
+        # Scatter - orbiting points with color gradient
+        n_points = 50
+        theta = np.linspace(0, 2 * np.pi, n_points)
+        x_scatter = 0.5 + 0.3 * np.cos(theta)
+        y_scatter = 0.5 + 0.3 * np.sin(theta)
+        colors = np.linspace(0, 1, n_points)
+        scatter = self.axes[4].scatter(
+            x_scatter, y_scatter, c=colors, cmap="plasma", s=100, alpha=0.8
+        )
+        self.axes[4].set_xlim(0, 1)
+        self.axes[4].set_ylim(0, 1)
+        self.axes[4].set_title("Orbiting Points")
+        self.axes[4].set_aspect("equal")
+        self.artists["scatter"] = scatter
+        self.artists["n_points"] = n_points
+
+        self.fig.tight_layout()
+        return self.fig
+
+    def update(self, frame_idx, params):
+        phase = params["phase"]
+
+        # Update line chart - multiple frequency sine wave
+        x = self.artists["line_x"]
+        y = np.sin(x + phase) + 0.3 * np.sin(3 * x - phase)
+        self.artists["line"].set_ydata(y)
+
+        # Update bar chart - pulsating bars
+        heights = np.abs(np.sin(np.arange(10) * 0.5 + phase)) * 1.5 + 0.2
+        for bar, h in zip(self.artists["bars"], heights):
+            bar.set_height(h)
+
+        # Update imshow - wave interference from moving sources
+        x = np.linspace(-5, 5, 30)
+        y = np.linspace(-5, 5, 30)
+        X, Y = np.meshgrid(x, y)
+
+        # Two moving wave sources
+        x1, y1 = 2 * np.cos(phase), 2 * np.sin(phase)
+        x2, y2 = 2 * np.cos(phase + np.pi), 2 * np.sin(phase + np.pi)
+
+        R1 = np.sqrt((X - x1) ** 2 + (Y - y1) ** 2)
+        R2 = np.sqrt((X - x2) ** 2 + (Y - y2) ** 2)
+
+        Z = np.sin(3 * R1 - 2 * phase) / (R1 + 0.5) + np.sin(3 * R2 - 2 * phase) / (
+            R2 + 0.5
+        )
+        Z = Z + 1  # Shift to positive range
+
+        self.artists["imshow"].set_data(Z)
+
+        # Update 3D line - rotating and pulsating helix
+        t = self.artists["line3d_t"]
+        radius = 1 + 0.3 * np.sin(2 * phase)
+        x_3d = radius * np.cos(t + phase)
+        y_3d = radius * np.sin(t + phase)
+        z_3d = t
+        self.artists["line3d"].set_data_3d(x_3d, y_3d, z_3d)
+
+        # Update scatter - orbiting and pulsating points
+        n_points = self.artists["n_points"]
+        angle = phase
+        r = 0.25 + 0.15 * np.sin(2 * angle)
+        theta = np.linspace(0, 2 * np.pi, n_points) + angle
+        x_scatter = 0.5 + r * np.cos(theta)
+        y_scatter = 0.5 + r * np.sin(theta)
+        self.artists["scatter"].set_offsets(np.c_[x_scatter, y_scatter])
+
+
+if __name__ == "__main__":
+    # Generate parameters
+    num_frames = 90
+    params = [{"phase": 2 * np.pi * i / num_frames} for i in range(num_frames)]
+
+    # Create animation
+    anim = MultiPanelAnimation()
+    output_path = Path("example_output/multi_panel_animation.mp4")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    anim.make_video(
+        output_file=output_path, param_by_frame=params, fps=30, num_workers=8
+    )

parallel_animate/examples/scaling_test.py

@@ -0,0 +1,96 @@
+import time
+import json
+from pathlib import Path
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+from parallel_animate.examples.very_complex_animation import VeryComplexAnimation
+
+
+def run_scaling_test(
+    num_frames: int, num_workers_list: list[int], output_dir: Path
+) -> dict[str, dict]:
+    """Run strong scaling test for different parallelization strategies."""
+    # Generate parameters
+    params = [{"phase": 2 * np.pi * i / num_frames} for i in range(num_frames)]
+
+    # Run tests for all configurations
+    results_all = {}
+    for num_workers in num_workers_list:
+        for reuse_figure_object in [True, False]:
+            config_name = f"{num_workers}workers_reusefig{reuse_figure_object}"
+            print(f"Running test: {config_name}...")
+            output_file = output_dir / f"output_{config_name}.mp4"
+            start_time = time.time()
+            anim = VeryComplexAnimation()
+            anim.make_video(
+                output_file=output_file,
+                param_by_frame=params,
+                fps=30,
+                num_workers=num_workers,
+                disable_progress_bar=True,
+                reuse_figure_object=reuse_figure_object,
+            )
+            elapsed_time = time.time() - start_time
+            results_all[config_name] = {
+                "num_workers": num_workers,
+                "reuse_fig_obj": reuse_figure_object,
+                "time_seconds": elapsed_time,
+            }
+            print(f"Test {config_name} completed in {elapsed_time:.2f} seconds")
+
+    return results_all
+
+
+def plot_scaling_results(results: dict[str, dict], output_path: Path) -> None:
+    """Create plots showing scaling performance."""
+    num_workers_list = list(set(result["num_workers"] for result in results.values()))
+
+    fig, ax = plt.subplots(figsize=(4, 4), tight_layout=True)
+    baseline_time = results[f"1workers_reusefigFalse"]["time_seconds"]
+    _max_speedup = len(num_workers_list)
+    for reuse_figure_object in [True, False]:
+        times = [
+            results[f"{n}workers_reusefig{reuse_figure_object}"]["time_seconds"]
+            for n in num_workers_list
+        ]
+        speedups = [baseline_time / t for t in times]
+        _max_speedup = max(_max_speedup, *speedups)
+        ax.plot(
+            num_workers_list,
+            speedups,
+            marker="o",
+            label="with cache" if reuse_figure_object else "without cache",
+        )
+    ax.plot(
+        [1, max(num_workers_list)],
+        [1, max(num_workers_list)],
+        color="black",
+        label="ideal scaling",
+        zorder=0,
+    )
+    ax.legend()
+    ax.set_xscale("log", base=2)
+    ax.set_yscale("log", base=2)
+    ax.set_xlabel("# workers")
+    ax.set_ylabel("speedup")
+    ax.set_title("Strong scaling test")
+    ax.set_aspect("equal")
+    fig.savefig(output_path)
+    plt.close(fig)
+
+
+if __name__ == "__main__":
+    output_dir = Path("example_output/scaling_test")
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    num_frames_to_draw = 320
+    num_workers_to_test = [1, 2, 4, 8, 16]
+    results = run_scaling_test(num_frames_to_draw, num_workers_to_test, output_dir)
+    with open(output_dir / "results.json", "w") as f:
+        json.dump(results, f, indent=4)
+
+    with open(output_dir / "results.json", "r") as f:
+        results = json.load(f)
+    plot_scaling_results(results, output_path=output_dir / "scaling_graph.png")

parallel_animate/examples/simple_wave_animation.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+from parallel_animate import Animator
+
+
+class WaveAnimation(Animator):
+    def setup(self):
+        fig, ax = plt.subplots()
+        self.x = np.linspace(0, 4 * np.pi, 200)
+        (self.line,) = ax.plot(self.x, np.cos(self.x))
+        ax.set_xlim(0, 4 * np.pi)
+        ax.set_ylim(-1.5, 1.5)
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        ax.set_title("Cosine Wave")
+        return fig
+
+    def update(self, frame_idx, params):
+        phase = params["phase"]
+        self.line.set_ydata(np.cos(self.x + phase))
+
+
+if __name__ == "__main__":
+    # Generate parameters
+    num_frames = 60
+    params = [{"phase": 2 * np.pi * i / num_frames} for i in range(num_frames)]
+
+    # Create and render
+    anim = WaveAnimation()
+    output_path = Path("example_output/simple_wave_animation.mp4")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    anim.make_video(
+        output_file=output_path, param_by_frame=params, fps=30, num_workers=4
+    )