diff --git a/src/data_morph/morpher.py b/src/data_morph/morpher.py index 3ad083e0..c2279a70 100644 --- a/src/data_morph/morpher.py +++ b/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 @@ def _record_frames( 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 @@ def _record_frames( 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 @@ def morph( 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 @@ def _easing( 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,26 +486,26 @@ def _easing( 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, @@ -547,7 +513,8 @@ def _easing( 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, ) diff --git a/src/data_morph/plotting/animation.py b/src/data_morph/plotting/animation.py index 7db01dcc..f3d12d80 100644 --- a/src/data_morph/plotting/animation.py +++ b/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, ) diff --git a/tests/plotting/test_animation.py b/tests/plotting/test_animation.py index 14c2980a..bb873062 100644 --- a/tests/plotting/test_animation.py +++ b/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'), diff --git a/tests/test_morpher.py b/tests/test_morpher.py index a3aaa4b0..b4892191 100644 --- a/tests/test_morpher.py +++ b/tests/test_morpher.py @@ -1,7 +1,6 @@ """Test the data_morph.morpher module.""" -import hashlib -from collections import Counter +import re from functools import partial import pandas as pd @@ -181,7 +180,7 @@ def test_no_writing(self, capsys): def test_saving_data(self, tmp_path): """Test that writing files to disk in the morph() method is working.""" - num_frames = 20 + num_frames = 5 iterations = 10 start_shape = 'dino' target_shape = 'circle' @@ -210,8 +209,12 @@ def test_saving_data(self, tmp_path): freeze_for=0, ) + frame_number_format = f'{{:0{len(str(iterations))}d}}'.format + # we don't save the data for the first frame since it is in the input data - assert not (tmp_path / f'{base_file_name}-data-000.csv').is_file() + assert not ( + tmp_path / f'{base_file_name}-data-{frame_number_format(0)}.csv' + ).is_file() # make sure we have the correct number of files for kind in ['png', 'csv']: @@ -219,7 +222,10 @@ def test_saving_data(self, tmp_path): # at the final frame, we have the output data assert_frame_equal( - pd.read_csv(tmp_path / f'{base_file_name}-data-{num_frames - 1:03d}.csv'), + pd.read_csv( + tmp_path + / f'{base_file_name}-data-{frame_number_format(iterations)}.csv' + ), morphed_data, ) @@ -227,7 +233,7 @@ def test_saving_data(self, tmp_path): with pytest.raises(AssertionError): assert_frame_equal( pd.read_csv( - tmp_path / f'{base_file_name}-data-{num_frames // 2:03d}.csv' + tmp_path / f'{base_file_name}-data-{frame_number_format(8)}.csv' ), morphed_data, ) @@ -237,16 +243,13 @@ def test_saving_data(self, tmp_path): @pytest.mark.parametrize('write_images', [True, False]) @pytest.mark.parametrize('start_frame', [0, 1, 20]) - @pytest.mark.parametrize('freeze_for', [0, 2, 10]) - def test_freeze_animation_frames( - self, write_images, start_frame, freeze_for, tmp_path - ): - """Confirm that freezing frames in the animation is working.""" + def test_record_frames(self, write_images, start_frame, tmp_path): + """Confirm that _record_frames() is working.""" dataset = DataLoader.load_dataset('dino') morpher = DataMorpher( decimals=2, write_images=write_images, - write_data=True, + write_data=False, output_dir=tmp_path, seed=21, keep_frames=True, @@ -254,30 +257,18 @@ def test_freeze_animation_frames( in_notebook=False, ) + frame_number = f'{start_frame:03d}' + base_path = 'test-freeze' - end_frame = morpher._record_frames( + morpher._record_frames( dataset.data, dataset.plot_bounds, base_path, - freeze_for, - start_frame, + frame_number, ) - # get image hashes - image_hashes = Counter() - for frame_image in tmp_path.glob(f'{base_path}*.png'): - with frame_image.open('rb') as img: - image_hashes.update({hashlib.sha256(img.read()).hexdigest(): 1}) - - if not write_images: - assert not image_hashes - else: - # check that the number of frames is correct - assert end_frame - start_frame == freeze_for - - # check that the images are indeed the same - if write_images and freeze_for: - assert len(image_hashes.keys()) == 1 - assert next(iter(image_hashes.values())) == freeze_for - else: - assert not image_hashes + images = list(tmp_path.glob(f'{base_path}*.png')) + + assert len(images) == int(write_images) + if write_images: + assert re.search(r'\d{3}', images[0].stem).group(0) == frame_number