Save only unique frames and use GIF frame durations from PIL

src/data_morph/morpher.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import shutil
 from contextlib import nullcontext
 from functools import partial
 from numbers import Number
@@ -200,9 +199,8 @@
         data: pd.DataFrame,
         bounds: BoundingBox,
         base_file_name: str,
-        count: int,
-        frame_number: int,
-    ) -> int:
+        frame_number: str,
+    ) -> None:
         """
         Record frame data as a plot and, when :attr:`write_data` is ``True``, as a CSV file.
 
@@ -214,55 +212,25 @@
             The plotting limits.
         base_file_name : str
             The prefix to the file names for both the PNG and GIF files.
-        count : int
-            The number of frames to record with the data.
-        frame_number : int
-            The starting frame number.
-
-        Returns
-        -------
-        int
-            The next frame number available for recording.
+        frame_number : str
+            The frame number with padding zeros added already.
         """
-        if (self.write_images or self.write_data) and count > 0:
-            is_start = frame_number == 0
-            img_file = (
-                self.output_dir / f'{base_file_name}-image-{frame_number:03d}.png'
+        if self.write_images:
+            plot(
+                data,
+                save_to=self.output_dir / f'{base_file_name}-image-{frame_number}.png',
+                decimals=self.decimals,
+                x_bounds=bounds.x_bounds,
+                y_bounds=bounds.y_bounds,
+                dpi=150,
             )
-            data_file = (
-                self.output_dir / f'{base_file_name}-data-{frame_number:03d}.csv'
+        if (
+            self.write_data and int(frame_number) > 0
+        ):  # don't write data for the initial frame (input data)
+            data.to_csv(
+                self.output_dir / f'{base_file_name}-data-{frame_number}.csv',
+                index=False,
             )
-            if self.write_images:
-                plot(
-                    data,
-                    save_to=img_file,
-                    decimals=self.decimals,
-                    x_bounds=bounds.x_bounds,
-                    y_bounds=bounds.y_bounds,
-                    dpi=150,
-                )
-            if (
-                self.write_data and not is_start
-            ):  # don't write data for the initial frame (input data)
-                data.to_csv(data_file, index=False)
-            frame_number += 1
-            if (duplicate_count := count - 1) > 0:
-                for _ in range(duplicate_count):
-                    if data_file.exists():
-                        shutil.copy(
-                            data_file,
-                            self.output_dir
-                            / f'{base_file_name}-data-{frame_number:03d}.csv',
-                        )
-                    if img_file.exists():
-                        shutil.copy(
-                            img_file,
-                            self.output_dir
-                            / f'{base_file_name}-image-{frame_number:03d}.png',
-                        )
-                    frame_number += 1
-
-        return frame_number
 
     def _is_close_enough(self, df1: pd.DataFrame, df2: pd.DataFrame) -> bool:
         """
@@ -477,11 +445,9 @@
             base_file_name=base_file_name,
             bounds=start_shape.plot_bounds,
         )
-        frame_number = record_frames(
-            data=morphed_data,
-            count=max(freeze_for, 1),
-            frame_number=0,
-        )
+
+        frame_number_format = f'{{:0{len(str(iterations))}d}}'.format
+        record_frames(data=morphed_data, frame_number=frame_number_format(0))
 
         def _easing(
             frame: int, *, min_value: Number, max_value: Number
@@ -507,7 +473,7 @@
                 task_id = progress_tracker.add_task(
                     f'{start_shape.name} to {target_shape}'
                 )
-            for i in range(iterations):
+            for i in range(1, iterations + 1):
                 perturbed_data = self._perturb(
                     morphed_data.copy(),
                     target_shape=target_shape,
@@ -520,34 +486,35 @@
                 if self._is_close_enough(start_shape.data, perturbed_data):
                     morphed_data = perturbed_data
 
-                frame_number = record_frames(
-                    data=morphed_data,
-                    count=frame_numbers.count(i),
-                    frame_number=frame_number,
-                )
+                if frame_numbers.count(i):
+                    record_frames(
+                        data=morphed_data, frame_number=frame_number_format(i)
+                    )
 
                 if progress_tracker:
                     progress_tracker.update(
                         task_id,
                         total=iterations,
-                        completed=i + 1,
-                        refresh=self._in_notebook and (i + 1) % 500 == 0,
+                        completed=i,
+                        refresh=self._in_notebook and (i) % 500 == 0,
                     )
                 else:
-                    progress[task_id] = {'progress': i + 1, 'total': iterations}
+                    progress[task_id] = {'progress': i, 'total': iterations}
 
         if self.write_images:
             stitch_gif_animation(
                 self.output_dir,
                 start_shape.name,
+                frame_numbers=frame_numbers,
                 target_shape=target_shape,
                 keep_frames=self.keep_frames,
                 forward_only_animation=self.forward_only_animation,
             )
 
         if self.write_data:
             morphed_data.to_csv(
-                self.output_dir / f'{base_file_name}-data-{frame_number:03d}.csv',
+                self.output_dir
+                / f'{base_file_name}-data-{frame_number_format(iterations)}.csv',
                 index=False,
             )
 

src/data_morph/plotting/animation.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import math
+import re
 from functools import wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
@@ -17,6 +18,7 @@ def stitch_gif_animation(
     output_dir: str | Path,
     start_shape: str,
     target_shape: str | Shape,
+    frame_numbers: list[int],
     keep_frames: bool = False,
     forward_only_animation: bool = False,
 ) -> None:
@@ -32,6 +34,10 @@ def stitch_gif_animation(
         The starting shape.
     target_shape : str or Shape
         The target shape for the morphing.
+    frame_numbers : list[int]
+        The saved frames to use in the GIF. Repeated consecutive frames will be shown
+        as a single frame for a longer duration (i.e., x repeats, means x times longer
+        than the default duration of 5 milliseconds).
     keep_frames : bool, default ``False``
         Whether to keep the individual frames after creating the animation.
     forward_only_animation : bool, default ``False``
@@ -44,22 +50,30 @@ def stitch_gif_animation(
         Frames are stitched together with Pillow.
     """
     output_dir = Path(output_dir)
+    iteration_pattern = re.compile(r'\d+')
+    default_frame_duration = 5  # milliseconds
 
     # find the frames and sort them
     imgs = sorted(output_dir.glob(f'{start_shape}-to-{target_shape}*.png'))
 
-    frames = [Image.open(img) for img in imgs]
+    frames = []
+    durations = []
+    for img_file in imgs:
+        iteration_number = int(iteration_pattern.search(img_file.stem).group(0))
+        frames.append(Image.open(img_file))
+        durations.append(frame_numbers.count(iteration_number) * default_frame_duration)
 
     if not forward_only_animation:
         # add the animation in reverse
         frames.extend(frames[::-1])
+        durations.extend(durations[::-1])
 
     frames[0].save(
         output_dir / f'{start_shape}_to_{target_shape}.gif',
         format='GIF',
         append_images=frames[1:],
         save_all=True,
-        duration=5,
+        duration=durations,
         loop=0,
     )
 

tests/plotting/test_animation.py

@@ -1,7 +1,10 @@
 """Test the animation module."""
 
+from contextlib import suppress
+
 import numpy as np
 import pytest
+from PIL import Image
 
 from data_morph.plotting import animation
 from data_morph.plotting.animation import stitch_gif_animation
@@ -10,14 +13,16 @@
 pytestmark = pytest.mark.plotting
 
 
-def test_frame_stitching(sample_data, tmp_path):
+@pytest.mark.parametrize('forward_only', [True, False])
+def test_frame_stitching(sample_data, tmp_path, forward_only):
     """Test stitching frames into a GIF animation."""
     start_shape = 'sample'
     target_shape = 'circle'
     bounds = [-5, 105]
+    frame_numbers = list(range(10))
     rng = np.random.default_rng()
 
-    for frame in range(10):
+    for frame in frame_numbers:
         plot(
             data=sample_data + rng.standard_normal(),
             x_bounds=bounds,
@@ -26,18 +31,44 @@ def test_frame_stitching(sample_data, tmp_path):
             decimals=2,
         )
 
+    duration_multipliers = [0, 0, 0, 0, 1, 1, *frame_numbers[2:], frame_numbers[-1]]
     stitch_gif_animation(
         output_dir=tmp_path,
         start_shape=start_shape,
         target_shape=target_shape,
+        frame_numbers=duration_multipliers,
         keep_frames=False,
-        forward_only_animation=False,
+        forward_only_animation=forward_only,
     )
 
     animation_file = tmp_path / f'{start_shape}_to_{target_shape}.gif'
     assert animation_file.is_file()
     assert not (tmp_path / f'{start_shape}-to-{target_shape}-{frame}.png').is_file()
 
+    with Image.open(animation_file) as img:
+        # we subtract one when playing in reverse as well because the middle frame (last
+        # in the forward direction) is combined into a single frame with the start of the
+        # reversal as part of PIL's optimization
+        assert img.n_frames == (
+            len(frame_numbers) if forward_only else len(frame_numbers) * 2 - 1
+        )
+        for frame in range(len(frame_numbers)):
+            with suppress(KeyError):
+                # if we play in reverse, the midpoint will have double duration since
+                # those two frames are combined
+                rewind_multiplier = (
+                    2 if not forward_only and frame == len(frame_numbers) - 1 else 1
+                )
+                # duration only seems to be present on frames where it is different
+                if frame_duration := img.info['duration']:
+                    assert (
+                        frame_duration
+                        == duration_multipliers.count(frame) * 5 * rewind_multiplier
+                    )
+            with suppress(EOFError):
+                # move to the next frame
+                img.seek(img.tell() + 1)
+
 
 @pytest.mark.parametrize(
     ('ease_function', 'step', 'expected'),