Fix CubeMove for rotated cubes

manim_rubikscube/cube.py

@@ -5,33 +5,35 @@
 from .cubie import Cubie
 from kociemba import solver as sv
 
-class RubiksCube(VMobject):
-    #If facing the Rubik's Cube, X goes Front to Back, Y goes Right to Left, Z goes Down to Up 
-    #Each coordinate starts at 0 and goes to (Dimensions - 1)
 
-    cubies = np.ndarray
-    indices = {}
+class RubiksCube(VMobject):
+    # If facing the Rubik's Cube, X goes Front to Back, Y goes Right to Left, Z goes Down to Up
+    # Each coordinate starts at 0 and goes to (Dimensions - 1)
 
     # Colors are in the order Up, Right, Front, Down, Left, Back
-    def __init__(self, dim=3, colors=[WHITE, "#B90000", "#009B48", "#FFD500", "#FF5900", "#0045AD"], x_offset=2.1, y_offset=2.1, z_offset=2.1):#, **kwargs):
+    def __init__(self, dim=3, colors=None, x_offset=2.1, y_offset=2.1, z_offset=2.1):
         if not (dim >= 2):
             raise Exception("Dimension must be >= 2")
 
         VMobject.__init__(self)
+
+        if colors is None:
+            colors = [WHITE, "#B90000", "#009B48", "#FFD500", "#FF5900", "#0045AD"]
+
         self.dimensions = dim
         self.colors = colors
         self.x_offset = [[Mobject.shift, [x_offset, 0, 0]]]
         self.y_offset = [[Mobject.shift, [0, y_offset, 0]]]
         self.z_offset = [[Mobject.shift, [0, 0, z_offset]]]
 
         self.cubies = np.ndarray((dim, dim, dim), dtype=Cubie)
-        self.generate_cubies()#**kwargs)
-    
-    def generate_cubies(self):#, **kwargs):
+        self.generate_cubies()
+
+    def generate_cubies(self):
         for x in range(self.dimensions):
             for y in range(self.dimensions):
                 for z in range(self.dimensions):
-                    cubie = Cubie(x, y, z, self.dimensions, self.colors)#, **kwargs)
+                    cubie = Cubie(x, y, z, self.dimensions, self.colors)
                     self.transform_cubie(x, self.x_offset, cubie)
                     self.transform_cubie(y, self.y_offset, cubie)
                     self.transform_cubie(z, self.z_offset, cubie)
@@ -55,20 +57,20 @@ def set_state(self, positions):
 
         for cubie in np.rot90(np.flip(self.get_face("R", False), (0, 1)), -1).flatten():
             cubie.get_face("R").set_fill(colors[positions.pop(0)], 1)
-        
+
         for cubie in np.rot90(np.flip(self.get_face("F", False), 0)).flatten():
             cubie.get_face("F").set_fill(colors[positions.pop(0)], 1)
-        
+
         for cubie in np.rot90(np.flip(self.get_face("D", False), 0), 2).flatten():
             cubie.get_face("D").set_fill(colors[positions.pop(0)], 1)
-        
+
         for cubie in np.rot90(np.flip(self.get_face("L", False), 0)).flatten():
             cubie.get_face("L").set_fill(colors[positions.pop(0)], 1)
-        
+
         for cubie in np.rot90(np.flip(self.get_face("B", False), (0, 1)), -1).flatten():
             cubie.get_face("B").set_fill(colors[positions.pop(0)], 1)
         # except:
-            # return
+        #     return
 
     def solve_by_kociemba(self, state):
         return sv.solve(state).replace("3", "'").replace("1", "").split()
@@ -88,31 +90,30 @@ def transform_cubie(self, position, offset, tile):
                         tile, magnitude * np.array(offset[i][1 + j * 2])
                     )
 
+    def get_face_slice(self, face):
+        """
+        Return a NumPy slice object specifying which part of the array corresponds
+        to which face. NumPy sli indexing a ndarray,
+        e.g. a[:, 2] == a[np.s_[:, 2]]
+        """
+        face_slices = {
+            "F": np.s_[0, :, :],
+            "B": np.s_[self.dimensions - 1, :, :],
+            "U": np.s_[:, :, self.dimensions - 1],
+            "D": np.s_[:, :, 0],
+            "L": np.s_[:, self.dimensions - 1, :],
+            "R": np.s_[:, 0, :],
+        }
+
+        if face in face_slices:
+            return face_slices[face]
+        else:
+            raise ValueError("Invalid face identifier " + face)
+
     def get_face(self, face, flatten=True):
-        if face == "F":
-            face = self.cubies[0, :, :]
-        elif face == "B":
-            face = self.cubies[self.dimensions-1, :, :]
-        elif face == "U":
-            face = self.cubies[:, :, self.dimensions-1]
-        elif face == "D":
-            face = self.cubies[:, :, 0]
-        elif face == "L":
-            face = self.cubies[:, self.dimensions-1, :]
-        elif face == "R":
-            face = self.cubies[:, 0, :]
+        face = self.cubies[self.get_face_slice(face)]
 
         if flatten:
             return face.flatten()
         else:
             return face
-
-    def set_indices(self):
-        for c in self.cubies.flatten():
-            # self.indices[c.get_position()] = c.get_position()
-            self.indices[c.get_rounded_center()] = c.position
-
-    def adjust_indices(self, cubies):
-        for c in cubies.flatten():
-            loc = self.indices[c.get_rounded_center()]
-            self.cubies[loc[0], loc[1], loc[2]] = c

manim_rubikscube/cube_animations.py

@@ -1,31 +1,51 @@
+import numpy as np
 from manim.animation.animation import Animation
 from manim.constants import PI
 from manim.mobject.types.vectorized_mobject import VGroup
-from .cube_utils import get_axis_from_face
+
+from .cube import RubiksCube
 
 class CubeMove(Animation):
-    def __init__(self, mobject, face, **kwargs):
-        self.axis = get_axis_from_face(face[0])
+    def __init__(self, mobject: RubiksCube, face, **kwargs):
+        # This only makes sense when called on a RubiksCube
+        assert isinstance(mobject, RubiksCube)
+
+        # Compute the axis of rotation by taking the vector from the cube's center
+        # to the middle cubie of the rotated face
+        # TODO: this might accumulate numerical errors, but it seems ok for tens of moves
+        self.axis = (
+            mobject.get_face(face[0], flatten=False)[1, 1].get_center()
+            - mobject.get_center()
+        )
         self.face = face
-        self.angle = PI/2 if ("R" in face or "F" in face or "D" in face) else -PI/2
-        self.angle = self.angle if "2" not in face else self.angle*2
-        self.angle = -self.angle if "'" in face else self.angle
+
+        self.n_turns = 1 if "2" not in face else 2
+        self.n_turns = -self.n_turns if "'" in face else self.n_turns
+
         super().__init__(mobject, **kwargs)
 
     def create_starting_mobject(self):
         starting_mobject = self.mobject.copy()
-        if starting_mobject.indices == {}:
-            starting_mobject.set_indices()
         return starting_mobject
 
     def interpolate_mobject(self, alpha):
         self.mobject.become(self.starting_mobject)
 
         VGroup(*self.mobject.get_face(self.face[0])).rotate(
-            alpha * self.angle,
+            -alpha * (PI / 2) * self.n_turns,
             self.axis
         )
 
     def finish(self):
         super().finish()
-        self.mobject.adjust_indices(self.mobject.get_face(self.face[0], False))
+        cubies = self.mobject.cubies[self.mobject.get_face_slice(self.face[0])]
+
+        # We need to make sure that moves that are supposed to be clockwise really are
+        n_turns = self.n_turns if (self.face[0] in {"L", "F", "D"}) else -self.n_turns
+
+        # Get to a non-negative value
+        n_turns = (n_turns + 4) % 4
+
+        cubies = np.rot90(cubies, k=n_turns)
+
+        self.mobject.cubies[self.mobject.get_face_slice(self.face[0])] = cubies

manim_rubikscube/cube_utils.py

@@ -1,16 +1,5 @@
 from manim.constants import *
 
-def get_axis_from_face(face):
-    if face == "F" or face == "B":
-        return X_AXIS
-    elif face == "U" or face == "D":
-        return Z_AXIS
-    else:
-        return Y_AXIS
-
-def get_direction_from_face(face):
-    #Clockwise/counterclockwise for each face. UP goes cw?
-    return
 
 def get_type_of_cubie(dim, position):
     if (position[1] == 0 or position[1] == dim-1) and (position[2] == 0 or position[2] == dim-1):

manim_rubikscube/cubie.py

@@ -32,7 +32,6 @@ def get_position(self):
     #     self.position = position
 
     def get_rounded_center(self):
-        #TODO: Switch from using center to cubie positions
         return tuple([round(self.get_x(), 3), round(self.get_y(), 3), round(self.get_z(), 3)])
 
     def generate_points(self):
@@ -84,4 +83,7 @@ def init_colors(self):
             width=self.background_stroke_width,
             opacity=self.background_stroke_opacity,
             family=False
-        )
+        )
+
+    def __repr__(self):
+        return "Cubie({})".format(self.get_rounded_center())