interface for shuffling and rotating a cube

README.md

@@ -12,7 +12,9 @@ uv sync
 pre-commit install
 ```
 
-## Basic usage
+## Usage
+
+### Create a cube
 
 ```python
 from rubik.cube import Cube
@@ -30,6 +32,20 @@ print(cube)
 #     DDD
 ```
 
+### Perform basic moves
+
+```python
+# shuffle the cube using 1000 random moves
+cube.shuffle(num_moves=1000, seed=0)
+print(cube)
+print(cube.history)
+
+# rotate it in some way
+cube.rotate('X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i')
+print(cube)
+print(cube.history)
+```
+
 ## Roadmap
 
 #### Fully tensorized Rubik Cube model

src/rubik/{moves.py → action.py}

@@ -1,11 +1,10 @@
+import numpy as np
 import torch
 import torch.nn.functional as F
 
-from rubik.cube import Cube
+from rubik.tensor_utils import build_permutation_matrix, build_cube_tensor
 
 
-INT8 = torch.int8
-
 POS_ROTATIONS = torch.stack(
     [
         # rot about X: Z -> Y
@@ -16,7 +15,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 0, 1],
                 [0, 0, -1, 0],
             ],
-            dtype=INT8,
+            dtype=torch.int8,
         ),
         # rot about Y: X -> Z
         torch.tensor(
@@ -26,7 +25,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 1, 0],
                 [0, 1, 0, 0],
             ],
-            dtype=INT8,
+            dtype=torch.int8,
         ),
         # rot about Z: Y -> X
         torch.tensor(
@@ -36,7 +35,7 @@ POS_ROTATIONS = torch.stack(
                 [0, -1, 0, 0],
                 [0, 0, 0, 1],
             ],
-            dtype=INT8,
+            dtype=torch.int8,
         ),
     ]
 )
@@ -47,49 +46,20 @@ POS_SHIFTS = torch.tensor(
         [0, 1, 0, 0],
         [0, 0, 1, 0],
     ],
-    dtype=INT8,
+    dtype=torch.int8,
 )
 
-FACE_PERMS = torch.stack(
+
+# rotation about X axis: 0 (Up)   -> 2 (Front) -> 5 (Down)  -> 4 (Back)  -> 0 (Up)
+# rotation about Y axis: 0 (Up)   -> 1 (Left)  -> 5 (Down)  -> 3 (Right) -> 0 (Up)
+# rotation about Z axis: 1 (Left) -> 2 (Front) -> 3 (Right) -> 4 (Back)  -> 1 (Left)
+FACE_ROTATIONS = torch.stack(
     [
-        # rotation about X axis: Up -> Front -> Down -> Back -> Up
-        torch.tensor(
-            [
-                [0, 0, 0, 0, 1, 0],
-                [0, 1, 0, 0, 0, 0],
-                [1, 0, 0, 0, 0, 0],
-                [0, 0, 0, 1, 0, 0],
-                [0, 0, 0, 0, 0, 1],
-                [0, 0, 1, 0, 0, 0],
-            ],
-            dtype=INT8,
-        ),
-        # rotation about Y axis: Up -> Left -> Down -> Right -> Up
-        torch.tensor(
-            [
-                [0, 0, 0, 1, 0, 0],
-                [1, 0, 0, 0, 0, 0],
-                [0, 0, 1, 0, 0, 0],
-                [0, 0, 0, 0, 0, 1],
-                [0, 0, 0, 0, 1, 0],
-                [0, 1, 0, 0, 0, 0],
-            ],
-            dtype=INT8,
-        ),
-        # rotation about Z axis: Left -> Front -> Right -> Back -> Left
-        torch.tensor(
-            [
-                [1, 0, 0, 0, 0, 0],
-                [0, 0, 0, 0, 1, 0],
-                [0, 1, 0, 0, 0, 0],
-                [0, 0, 1, 0, 0, 0],
-                [0, 0, 0, 1, 0, 0],
-                [0, 0, 0, 0, 0, 1],
-            ],
-            dtype=INT8,
-        ),
+        build_permutation_matrix(size=6, perm="0254"),
+        build_permutation_matrix(size=6, perm="0153"),
+        build_permutation_matrix(size=6, perm="1234"),
     ]
-).transpose(1, 2)
+)
 
 
 def build_actions_tensor(size: int) -> torch.Tensor:
@@ -98,29 +68,28 @@ def build_actions_tensor(size: int) -> torch.Tensor:
     """
     return torch.stack(
         [
-            build_permunation_tensor(size=size, axis=axis, slice=slice, inverse=inverse)
+            build_action_tensor(size=size, axis=axis, slice=slice, inverse=inverse)
             for axis in range(3)
             for slice in range(size)
             for inverse in range(2)
         ],
         dim=0,
-    ).sum(dim=0, dtype=INT8)
+    ).sum(dim=0, dtype=torch.int8)
 
 
-def build_permunation_tensor(size: int, axis: int, slice: int, inverse: int) -> torch.Tensor:
+def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch.Tensor:
     """
     Compute the sparse permutation tensor whose effect on a position-frozen color vector
     is the rotation along the specified axis, within the specified slice and the specified
     orientation.
     """
-    cube = Cube.create(["U", "L", "C", "R", "B", "D"], size=size)
+    tensor = build_cube_tensor(colors=list("ULCRBD"), size=size)
     length = 6 * (size**2)
 
     # extract faces impacted by the move
-    coordinates: torch.Tensor = cube.coordinates  # size = (length, 4)
-    transposed = coordinates.transpose(0, 1)  # size = (4, length)
-    indices = (transposed[axis + 1] == slice).nonzero().reshape(-1)  # size = (n,), n < length
-    extract = transposed[:, indices]  # size = (4, n)
+    indices = tensor.indices().to(dtype=torch.int8)  # size = (4, length)
+    changes = (indices[axis + 1] == slice).nonzero().reshape(-1)  # size = (n,), n < length
+    extract = indices[:, changes]  # size = (4, n)
 
     # apply coordinate rotation
     rotated = POS_ROTATIONS[axis] @ extract  # size = (4, n)
@@ -128,7 +97,7 @@ def build_permunation_tensor(size: int, axis: int, slice: int, inverse: int) ->
     rotated = rotated + offsets  # size = (4, n)
 
     # apply face rotation
-    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(INT8) @ FACE_PERMS[axis]).argmax(dim=-1)
+    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(torch.int8) @ FACE_ROTATIONS[axis]).argmax(dim=-1)
 
     # from this point on, convert rotation into a position-based permutation of colors
     (inputs, outputs) = (rotated, extract) if bool(inverse) else (extract, rotated)
@@ -136,7 +105,7 @@ def build_permunation_tensor(size: int, axis: int, slice: int, inverse: int) ->
     outputs = outputs.transpose(0, 1).tolist()  # size = (n, 4)
 
     # compute position-based permutation of colors equivalent to rotation converting inputs into outputs
-    local_to_total = dict(enumerate(indices.tolist()))
+    local_to_total = dict(enumerate(changes.tolist()))
     total_to_local = {ind: i for i, ind in local_to_total.items()}
 
     local_perm = {i: inputs.index(outputs[i]) for i in range(len(inputs))}
@@ -147,8 +116,29 @@ def build_permunation_tensor(size: int, axis: int, slice: int, inverse: int) ->
     # convert permutation dict into sparse tensor
     perm_indices = torch.tensor(
         [[axis] * length, [slice] * length, [inverse] * length, list(total_perm.keys()), list(total_perm.values())],
-        dtype=INT8,
+        dtype=torch.int8,
     )
-    perm_values = torch.tensor([1] * length, dtype=INT8)
+    perm_values = torch.tensor([1] * length, dtype=torch.int8)
     perm_size = (3, size, 2, length, length)
-    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=INT8)
+    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=torch.int8)
+
+
+def parse_action_str(name: str) -> tuple[int, ...]:
+    """
+    Convert the name of an action into a triple (axis, slice, inverse).
+    Examples:
+        'X1'  -> (0, 1, 0)
+        'X2i' -> (0, 2, 1)
+    """
+    return ("XYZ".index(name[0]), int(name[1]), int(len(name) >= 3))
+
+
+def sample_actions_str(num_moves: int, size: int, seed: int = 0) -> str:
+    """
+    Generate a string containing moves that are randomly sampled.
+    """
+    rng = np.random.default_rng(seed=seed)
+    axes = rng.choice(["X", "Y", "Z"], size=num_moves)
+    slices = rng.choice([str(i) for i in range(size)], size=num_moves)
+    orients = rng.choice(["", "i"], size=num_moves)
+    return " ".join("".join(move) for move in zip(axes, slices, orients))

src/rubik/cube.py

@@ -2,6 +2,10 @@ from dataclasses import dataclass
 
 import torch
 
+from rubik.action import build_actions_tensor, parse_action_str, sample_actions_str
+from rubik.display import stringify
+from rubik.tensor_utils import build_cube_tensor
+
 
 @dataclass
 class Cube:
@@ -18,6 +22,8 @@ class Cube:
 
     coordinates: torch.Tensor
     state: torch.Tensor
+    actions: torch.Tensor
+    history: list[list[int]]
     colors: list[str]
     size: int
 
@@ -28,66 +34,55 @@ class Cube:
         Example:
             cube = Cube.create(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
         """
-        assert (num := len(set(colors))) == 6, f"Expected 6 distinct colors, got {num}"
-        assert isinstance(size, int) and size > 1, f"Expected non-zero integrer size, got {size}"
-
-        # build dense tensor filled with colors
-        n = size - 1
-        tensor = torch.zeros([6, size, size, size], dtype=torch.int8)
-        tensor[0, :, :, n] = 1  # up
-        tensor[1, 0, :, :] = 2  # left
-        tensor[2, :, n, :] = 3  # front
-        tensor[3, n, :, :] = 4  # right
-        tensor[4, :, 0, :] = 5  # back
-        tensor[5, :, :, 0] = 6  # down
-        return cls.from_sparse(tensor.to_sparse(), colors, size)
-
-    def shuffle(self, num_moves: int):
-        raise NotImplementedError
+        tensor = build_cube_tensor(colors, size)
+        coordinates = tensor.indices().transpose(0, 1).to(torch.int8)
+        values = tensor.values()
+        actions = build_actions_tensor(size)
+        history: list[list[int]] = []
+        return cls(coordinates, values, actions, history, colors, size)
 
-    def rotate(self, moves: str):
-        raise NotImplementedError
+    def reset_history(self) -> None:
+        """
+        Reset internal history of moves.
+        """
+        self.history = []
+        return
 
-    def solve(slef, policy: str):
-        raise NotImplementedError
+    def shuffle(self, num_moves: int, seed: int = 0) -> None:
+        """
+        Randomly shuffle the cube by the supplied number of steps, and reset history of moves.
+        """
+        moves = sample_actions_str(num_moves, self.size, seed=seed)
+        self.rotate(moves)
+        self.reset_history()
+        return
 
-    @staticmethod
-    def pad_colors(colors: list[str]) -> list[str]:
+    def rotate(self, moves: str) -> None:
         """
-        Pad color names to strings of equal length.
+        Apply a sequence of moves (defined as plain string) to the cube.
         """
-        max_len = max(len(c) for c in colors)
-        return [c + " " * (max_len - len(c)) for c in colors]
+        actions = [parse_action_str(move) for move in moves.strip().split()]
+        for action in actions:
+            self.rotate_once(*action)
+        return
 
-    @classmethod
-    def from_sparse(cls, tensor: torch.Tensor, colors: list[str], size: int) -> "Cube":
+    def rotate_once(self, axis: int, slice: int, inverse: int) -> None:
         """
-        Gather cube attributes into a torch sparse tensor.
+        Apply a move (defined as 3 coordinates) to the cube.
         """
-        coordinates = tensor.indices().transpose(0, 1).to(torch.int8)
-        values = tensor.values()
-        return cls(coordinates, values, colors, size)
+        action = self.actions[axis, slice, inverse]
+        self.state = action @ self.state
+        self.history.append([axis, slice, inverse])
+        return
 
-    def to_sparse(self) -> torch.Tensor:
+    def solve(self, policy: str) -> None:
         """
-        Gather cube attributes into a torch sparse tensor.
+        Apply the specified solving policy to the cube.
         """
-        return torch.sparse_coo_tensor(
-            indices=self.coordinates.transpose(0, 1),
-            values=self.state,
-            size=(6, self.size, self.size, self.size),
-            dtype=torch.int8,
-        )
+        raise NotImplementedError
 
     def __str__(self):
         """
         Compute a string representation of a cube.
         """
-        colors = self.pad_colors(self.colors)
-        faces = self.state.reshape(6, self.size, self.size).transpose(1, 2)
-        faces = [[[colors[i - 1] for i in row] for row in face.tolist()] for face in faces]
-        void = " " * max(len(c) for c in self.colors) * self.size
-        l1 = "\n".join(" ".join([void, "".join(row), void, void]) for row in faces[0])
-        l2 = "\n".join(" ".join("".join(face[i]) for face in faces[1:5]) for i in range(self.size))
-        l3 = "\n".join(" ".join((void, "".join(row), void, void)) for row in faces[-1])
-        return "\n".join([l1, l2, l3])
+        return stringify(self)

src/rubik/display.py

@@ -0,0 +1,20 @@
+def stringify(cube) -> str:
+    """
+    Compute a string representation of a cube.
+    """
+    colors = pad_colors(cube.colors)
+    faces = cube.state.reshape(6, cube.size, cube.size).transpose(1, 2)
+    faces = [[[colors[i - 1] for i in row] for row in face.tolist()] for face in faces]
+    space = " " * max(len(c) for c in cube.colors) * cube.size
+    l1 = "\n".join(" ".join([space, "".join(row), space, space]) for row in faces[0])
+    l2 = "\n".join(" ".join("".join(face[i]) for face in faces[1:5]) for i in range(cube.size))
+    l3 = "\n".join(" ".join((space, "".join(row), space, space)) for row in faces[-1])
+    return "\n".join([l1, l2, l3])
+
+
+def pad_colors(colors: list[str]) -> list[str]:
+    """
+    Pad color names to strings of equal length.
+    """
+    max_len = max(len(c) for c in colors)
+    return [c + " " * (max_len - len(c)) for c in colors]

src/rubik/tensor_utils.py

@@ -0,0 +1,31 @@
+import torch
+
+
+def build_cube_tensor(colors: list[str], size: int) -> torch.Tensor:
+    """
+    Convert a list of 6 colors and size into a sparse 4D tensor representing a cube.
+    """
+    assert (num := len(set(colors))) == 6, f"Expected 6 distinct colors, got {num}"
+    assert isinstance(size, int) and size > 1, f"Expected non-zero integrer size, got {size}"
+
+    # build dense tensor filled with colors
+    n = size - 1
+    tensor = torch.zeros([6, size, size, size], dtype=torch.int8)
+    tensor[0, :, :, n] = 1  # up
+    tensor[1, 0, :, :] = 2  # left
+    tensor[2, :, n, :] = 3  # front
+    tensor[3, n, :, :] = 4  # right
+    tensor[4, :, 0, :] = 5  # back
+    tensor[5, :, :, 0] = 6  # down
+    return tensor.to_sparse()
+
+
+def build_permutation_matrix(size: int, perm: str) -> torch.Tensor:
+    """
+    Convert a permutation sting into a sparse 2D matrix.
+    """
+    perm_list = [int(p) for p in (perm + perm[0])]
+    perm_dict = {perm_list[i]: perm_list[i + 1] for i in range(len(perm))}
+    indices = torch.tensor([list(range(size)), [(perm_dict.get(i, i)) for i in range(size)]], dtype=torch.int8)
+    values = torch.tensor([1] * size, dtype=torch.int8)
+    return torch.sparse_coo_tensor(indices=indices, values=values, size=(size, size), dtype=torch.int8)