From 2d709758d094f4bdd29ad4eb00cdf283e23f4948 Mon Sep 17 00:00:00 2001
From: Grant Sanderson <grant@3blue1brown.com>
Date: Mon, 4 Jan 2021 23:09:29 -0800
Subject: [PATCH] Added TransformMatchingParts for smarter tex transforms

---
 .../animation/transform_matching_parts.py     | 51 +++++++++++++++++++
 manimlib/imports.py                           |  1 +
 2 files changed, 52 insertions(+)
 create mode 100644 manimlib/animation/transform_matching_parts.py

diff --git a/manimlib/animation/transform_matching_parts.py b/manimlib/animation/transform_matching_parts.py
new file mode 100644
index 00000000..1168c238
--- /dev/null
+++ b/manimlib/animation/transform_matching_parts.py
@@ -0,0 +1,51 @@
+from manimlib.animation.composition import AnimationGroup
+from manimlib.animation.fading import FadeTransformPieces
+from manimlib.animation.transform import Transform
+
+from manimlib.mobject.types.vectorized_mobject import VGroup
+from manimlib.mobject.types.vectorized_mobject import VMobject
+
+
+class TransformMatchingParts(AnimationGroup):
+    def __init__(self, mobject, target_mobject, **kwargs):
+        assert(isinstance(mobject, VMobject) and isinstance(target_mobject, VMobject))
+        source_map = self.get_shape_map(mobject)
+        target_map = self.get_shape_map(target_mobject)
+
+        transform_source = VGroup()
+        transform_target = VGroup()
+        fade_source = VGroup()
+        fade_target = VGroup()
+
+        for key in set(source_map).intersection(target_map):
+            transform_source.add(source_map[key])
+            transform_target.add(target_map[key])
+        for key in set(source_map).difference(target_map):
+            fade_source.add(source_map[key])
+        for key in set(target_map).difference(source_map):
+            fade_target.add(target_map[key])
+
+        kwargs["final_alpha_value"] = 0
+        super().__init__(
+            Transform(transform_source, transform_target, **kwargs),
+            FadeTransformPieces(fade_source, fade_target, **kwargs),
+        )
+
+        self.to_remove = mobject
+        self.to_add = target_mobject
+
+    def get_shape_map(self, mobject):
+        shape_map = {}
+        for sm in mobject.family_members_with_points():
+            key = hash(sm.get_triangulation().tobytes())
+            if key not in shape_map:
+                shape_map[key] = VGroup()
+            shape_map[key].add(sm)
+        return shape_map
+
+    def clean_up_from_scene(self, scene):
+        for anim in self.animations:
+            anim.update(0)
+        scene.remove(self.mobject)
+        scene.remove(self.to_remove)
+        scene.add(self.to_add)
diff --git a/manimlib/imports.py b/manimlib/imports.py
index 671327d5..83966a34 100644
--- a/manimlib/imports.py
+++ b/manimlib/imports.py
@@ -27,6 +27,7 @@ from manimlib.animation.numbers import *
 from manimlib.animation.rotation import *
 from manimlib.animation.specialized import *
 from manimlib.animation.transform import *
+from manimlib.animation.transform_matching_parts import *
 from manimlib.animation.update import *
 
 from manimlib.camera.camera import *