Added Scene caching

manim/default.cfg

@@ -97,6 +97,12 @@ frame_rate = 60
 pixel_height = 1440
 pixel_width = 2560
 
+# Use -1 to set max_files_cached to infinity.
+max_files_cached = 100
+#Flush cache will delete all the cached partial-movie-files.
+flush_cache = False
+disable_caching = False
+
 # These override the previous by using -t, --transparent
 [transparent]
 png_mode = RGBA

manim/scene/scene.py

@@ -2,6 +2,7 @@
 import random
 import warnings
 import platform
+import copy
 
 from tqdm import tqdm as ProgressDisplay
 import numpy as np
@@ -16,6 +17,7 @@
 from ..mobject.mobject import Mobject
 from ..scene.scene_file_writer import SceneFileWriter
 from ..utils.iterables import list_update
+from ..utils.hashing import get_hash_from_play_call, get_hash_from_wait_call
 
 
 class Scene(Container):
@@ -55,7 +57,7 @@ def __init__(self, **kwargs):
         Container.__init__(self, **kwargs)
         self.camera = self.camera_class(**camera_config)
         self.file_writer = SceneFileWriter(self, **file_writer_config,)
-
+        self.play_hashes_list = []
         self.mobjects = []
         # TODO, remove need for foreground mobjects
         self.foreground_mobjects = []
@@ -72,6 +74,9 @@ def __init__(self, **kwargs):
         except EndSceneEarlyException:
             pass
         self.tear_down()
+        # We have to reset these settings in case of multiple renders.
+        file_writer_config["skip_animations"] = False
+        self.original_skipping_status = file_writer_config["skip_animations"]
         self.file_writer.finish()
         self.print_end_message()
 
@@ -373,6 +378,17 @@ def add_mobjects_among(self, values):
         self.add(*filter(lambda m: isinstance(m, Mobject), values))
         return self
 
+    def add_mobjects_from_animations(self, animations):
+
+        curr_mobjects = self.get_mobject_family_members()
+        for animation in animations:
+            # Anything animated that's not already in the
+            # scene gets added to the scene
+            mob = animation.mobject
+            if mob not in curr_mobjects:
+                self.add(mob)
+                curr_mobjects += mob.get_family()
+
     def remove(self, *mobjects):
         """
         Removes mobjects in the passed list of mobjects
@@ -832,6 +848,71 @@ def update_skipping_status(self):
                 file_writer_config["skip_animations"] = True
                 raise EndSceneEarlyException()
 
+    def handle_caching_play(func):
+        """
+        Decorator that returns a wrapped version of func that will compute the hash of the play invocation.
+
+        The returned function will act according to the computed hash: either skip the animation because it's already cached, or let the invoked function play normally.
+
+        Parameters
+        ----------
+        func : Callable[[...], None]    
+            The play like function that has to be written to the video file stream. Take the same parameters as `scene.play`.
+        """
+
+        def wrapper(self, *args, **kwargs):
+            self.revert_to_original_skipping_status()
+            animations = self.compile_play_args_to_animation_list(*args, **kwargs)
+            self.add_mobjects_from_animations(animations)
+            if not file_writer_config["disable_caching"]:
+                mobjects_on_scene = self.get_mobjects()
+                hash_play = get_hash_from_play_call(
+                    self.camera, animations, mobjects_on_scene
+                )
+                self.play_hashes_list.append(hash_play)
+                if self.file_writer.is_already_cached(hash_play):
+                    logger.info(
+                        f"Animation {self.num_plays} : Using cached data (hash : {hash_play})"
+                    )
+                    file_writer_config["skip_animations"] = True
+            else:
+                hash_play = "uncached_{:05}".format(self.num_plays)
+                self.play_hashes_list.append(hash_play)
+            func(self, *args, **kwargs)
+
+        return wrapper
+
+    def handle_caching_wait(func):
+        """
+        Decorator that returns a wrapped version of func that will compute the hash of the wait invocation.
+
+        The returned function will act according to the computed hash: either skip the animation because it's already cached, or let the invoked function play normally.
+
+        Parameters
+        ----------
+        func : Callable[[...], None]
+            The wait like function that has to be written to the video file stream. Take the same parameters as `scene.wait`.
+        """
+
+        def wrapper(self, duration=DEFAULT_WAIT_TIME, stop_condition=None):
+            self.revert_to_original_skipping_status()
+            if not file_writer_config["disable_caching"]:
+                hash_wait = get_hash_from_wait_call(
+                    self.camera, duration, stop_condition, self.get_mobjects()
+                )
+                self.play_hashes_list.append(hash_wait)
+                if self.file_writer.is_already_cached(hash_wait):
+                    logger.info(
+                        f"Wait {self.num_plays} : Using cached data (hash : {hash_wait})"
+                    )
+                    file_writer_config["skip_animations"] = True
+            else:
+                hash_wait = "uncached_{:05}".format(self.num_plays)
+                self.play_hashes_list.append(hash_wait)
+            func(self, duration, stop_condition)
+
+        return wrapper
+
     def handle_play_like_call(func):
         """
         This method is used internally to wrap the
@@ -875,16 +956,9 @@ def begin_animations(self, animations):
             List of involved animations.
 
         """
-        curr_mobjects = self.get_mobject_family_members()
         for animation in animations:
             # Begin animation
             animation.begin()
-            # Anything animated that's not already in the
-            # scene gets added to the scene
-            mob = animation.mobject
-            if mob not in curr_mobjects:
-                self.add(mob)
-                curr_mobjects += mob.get_family()
 
     def progress_through_animations(self, animations):
         """
@@ -933,6 +1007,7 @@ def finish_animations(self, animations):
         else:
             self.update_mobjects(0)
 
+    @handle_caching_play
     @handle_play_like_call
     def play(self, *args, **kwargs):
         """
@@ -1032,6 +1107,7 @@ def get_wait_time_progression(self, duration, stop_condition):
             time_progression.set_description("Waiting {}".format(self.num_plays))
         return time_progression
 
+    @handle_caching_wait
     @handle_play_like_call
     def wait(self, duration=DEFAULT_WAIT_TIME, stop_condition=None):
         """
@@ -1105,8 +1181,8 @@ def force_skipping(self):
         Scene
             The Scene, with skipping turned on.
         """
-        self.original_skipping_status = self.SKIP_ANIMATIONS
-        self.SKIP_ANIMATIONS = True
+        self.original_skipping_status = file_writer_config["skip_animations"]
+        file_writer_config["skip_animations"] = True
         return self
 
     def revert_to_original_skipping_status(self):
@@ -1121,7 +1197,7 @@ def revert_to_original_skipping_status(self):
             The Scene, with the original skipping status.
         """
         if hasattr(self, "original_skipping_status"):
-            self.SKIP_ANIMATIONS = self.original_skipping_status
+            file_writer_config["skip_animations"] = self.original_skipping_status
         return self
 
     def add_frames(self, *frames):
@@ -1156,7 +1232,7 @@ def add_sound(self, sound_file, time_offset=0, gain=None, **kwargs):
         gain :
 
         """
-        if self.SKIP_ANIMATIONS:
+        if file_writer_config["skip_animations"]:
             return
         time = self.get_time() + time_offset
         self.file_writer.add_sound(sound_file, time, gain, **kwargs)

manim/scene/scene_file_writer.py

@@ -14,7 +14,7 @@
 from ..utils.config_ops import digest_config
 from ..utils.file_ops import guarantee_existence
 from ..utils.file_ops import add_extension_if_not_present
-from ..utils.file_ops import get_sorted_integer_files
+from ..utils.file_ops import modify_atime
 from ..utils.sounds import get_full_sound_file_path
 
 
@@ -170,8 +170,9 @@ def get_next_partial_movie_path(self):
         """
         result = os.path.join(
             self.partial_movie_directory,
-            "{:05}{}".format(
-                self.scene.num_plays, file_writer_config["movie_file_extension"],
+            "{}{}".format(
+                self.scene.play_hashes_list[self.scene.num_plays],
+                file_writer_config["movie_file_extension"],
             ),
         )
         return result
@@ -351,6 +352,10 @@ def finish(self):
             if hasattr(self, "writing_process"):
                 self.writing_process.terminate()
             self.combine_movie_files()
+            if file_writer_config["flush_cache"]:
+                self.flush_cache_directory()
+            else:
+                self.clean_cache()
         if file_writer_config["save_last_frame"]:
             self.scene.update_frame(ignore_skipping=True)
             self.save_final_image(self.scene.get_image())
@@ -421,6 +426,28 @@ def close_movie_pipe(self):
         shutil.move(
             self.temp_partial_movie_file_path, self.partial_movie_file_path,
         )
+        logger.debug(
+            f"Animation {self.scene.num_plays} : Partial movie file written in {self.partial_movie_file_path}"
+        )
+
+    def is_already_cached(self, hash_invocation):
+        """Will check if a file named with `hash_invocation` exists.
+
+        Parameters
+        ----------
+        hash_invocation : :class:`str`
+            The hash corresponding to an invocation to either `scene.play` or `scene.wait`.
+
+        Returns
+        -------
+        :class:`bool`
+            Whether the file exists.
+        """
+        path = os.path.join(
+            self.partial_movie_directory,
+            "{}{}".format(hash_invocation, self.movie_file_extension),
+        )
+        return os.path.exists(path)
 
     def combine_movie_files(self):
         """
@@ -435,34 +462,27 @@ def combine_movie_files(self):
         # cuts at all the places you might want.  But for viewing
         # the scene as a whole, one of course wants to see it as a
         # single piece.
-        kwargs = {
-            "remove_non_integer_files": True,
-            "extension": file_writer_config["movie_file_extension"],
-        }
-        if file_writer_config["from_animation_number"] is not None:
-            kwargs["min_index"] = file_writer_config["from_animation_number"]
-        if file_writer_config["upto_animation_number"] not in [None, np.inf]:
-            kwargs["max_index"] = file_writer_config["upto_animation_number"]
-        else:
-            kwargs["remove_indices_greater_than"] = self.scene.num_plays - 1
-        partial_movie_files = get_sorted_integer_files(
-            self.partial_movie_directory, **kwargs
-        )
+        partial_movie_files = [
+            os.path.join(
+                self.partial_movie_directory,
+                "{}{}".format(hash_play, file_writer_config["movie_file_extension"]),
+            )
+            for hash_play in self.scene.play_hashes_list
+        ]
         if len(partial_movie_files) == 0:
             logger.error("No animations in this scene")
             return
-
         # Write a file partial_file_list.txt containing all
-        # partial movie files
+        # partial movie files. This is used by FFMPEG.
         file_list = os.path.join(
             self.partial_movie_directory, "partial_movie_file_list.txt"
         )
         with open(file_list, "w") as fp:
+            fp.write("# This file is used internally by FFMPEG.\n")
             for pf_path in partial_movie_files:
                 if os.name == "nt":
                     pf_path = pf_path.replace("\\", "/")
                 fp.write("file 'file:{}'\n".format(pf_path))
-
         movie_file_path = self.get_movie_file_path()
         commands = [
             FFMPEG_BIN,
@@ -527,6 +547,46 @@ def combine_movie_files(self):
             os.remove(sound_file_path)
 
         self.print_file_ready_message(movie_file_path)
+        if file_writer_config["write_to_movie"]:
+            for file_path in partial_movie_files:
+                # We have to modify the accessed time so if we have to clean the cache we remove the one used the longest.
+                modify_atime(file_path)
+
+    def clean_cache(self):
+        """Will clean the cache by removing the partial_movie_files used by manim the longest ago."""
+        cached_partial_movies = [
+            os.path.join(self.partial_movie_directory, file_name)
+            for file_name in os.listdir(self.partial_movie_directory)
+            if file_name != "partial_movie_file_list.txt"
+        ]
+        if len(cached_partial_movies) > file_writer_config["max_files_cached"]:
+            number_files_to_delete = (
+                len(cached_partial_movies) - file_writer_config["max_files_cached"]
+            )
+            oldest_files_to_delete = sorted(
+                [partial_movie_file for partial_movie_file in cached_partial_movies],
+                key=os.path.getatime,
+            )[:number_files_to_delete]
+            # oldest_file_path = min(cached_partial_movies, key=os.path.getatime)
+            for file_to_delete in oldest_files_to_delete:
+                os.remove(file_to_delete)
+            logger.info(
+                f"The partial movie directory is full (> {file_writer_config['max_files_cached']} files). Therefore, manim has removed {number_files_to_delete} file(s) used by it the longest ago."
+                + "You can change this behaviour by changing max_files_cached in config."
+            )
+
+    def flush_cache_directory(self):
+        """Delete all the cached partial movie files"""
+        cached_partial_movies = [
+            os.path.join(self.partial_movie_directory, file_name)
+            for file_name in os.listdir(self.partial_movie_directory)
+            if file_name != "partial_movie_file_list.txt"
+        ]
+        for f in cached_partial_movies:
+            os.remove(f)
+        logger.info(
+            f"Cache flushed. {len(cached_partial_movies)} file(s) deleted in {self.partial_movie_directory}."
+        )
 
     def print_file_ready_message(self, file_path):
         """