Faster moves (#3)

* edit readme

* edit readme

* add docstring

* replace matmul by torch gathering for permutations of facelets

README.md

@@ -7,7 +7,7 @@ This project uses `uv 0.7` as environment & dependency manager, and `python 3.11
 
 ```shell
 uv venv
--- Activate env --
+(Activate env)
 uv sync
 pre-commit install
 ```
@@ -67,13 +67,11 @@ cube.rotate('X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i')
 
 #### Base solvers following rule-based policies
 
-
-
 ## References
 
-### Python implementations & rule-based solvers
+### Implementations & rule-based solvers
 
-Open-source projects related to Rubik's Cube, sorted by number of stars:
+Open-source projects related to Rubik's Cube:
 - [adrianliaw/PyCuber](https://github.com/adrianliaw/PyCuber)
 - [pglass/cube](https://github.com/pglass/cube)
 - [dwalton76/rubiks-cube-NxNxN-solver](https://github.com/dwalton76/rubiks-cube-NxNxN-solver)
@@ -81,6 +79,7 @@ Open-source projects related to Rubik's Cube, sorted by number of stars:
 - [trincaog/magiccube](https://github.com/trincaog/magiccube)
 - [charlstown/rubiks-cube-solver](https://github.com/charlstown/rubiks-cube-solver)
 - [staetyk/NxNxN-Cubes](https://github.com/staetyk/NxNxN-Cubes)
+- [wata-orz/santa2023_permutation_puzzle](https://github.com/wata-orz/santa2023_permutation_puzzle/tree/main)
 
 ### Machine Learning based solver models
 

notebooks/dev.ipynb

@@ -44,8 +44,6 @@
    "source": [
     "size = 3\n",
     "\n",
-    "actions = build_actions_tensor(size)\n",
-    "\n",
     "cube = Cube(size)\n",
     "print(cube)"
    ]
@@ -85,7 +83,7 @@
    "outputs": [],
    "source": [
     "cubis = copy.deepcopy(cube)\n",
-    "cubis.scramble(2000, seed=0)\n",
+    "cubis.scramble(20000, seed=0)\n",
     "print(cubis)\n",
     "print(cubis.history)"
    ]
@@ -110,7 +108,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "(actions[0, 2, 0].type(torch.float32) @ actions[0, 1, 1].type(torch.float32)).type(torch.int8)"
+    "actions = build_actions_tensor(size)\n",
+    "\n",
+    "torch.gather(actions[0, 2, 0], 0, actions[0, 1, 1])"
    ]
   }
  ],

src/rubik/action.py

@@ -17,7 +17,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 0, 1],
                 [0, 0, -1, 0],
             ],
-            dtype=torch.int32,
+            dtype=torch.int64,
         ),
         # rot about Y: X -> Z
         torch.tensor(
@@ -27,7 +27,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 1, 0],
                 [0, 1, 0, 0],
             ],
-            dtype=torch.int32,
+            dtype=torch.int64,
         ),
         # rot about Z: Y -> X
         torch.tensor(
@@ -37,7 +37,7 @@ POS_ROTATIONS = torch.stack(
                 [0, -1, 0, 0],
                 [0, 0, 0, 1],
             ],
-            dtype=torch.int32,
+            dtype=torch.int64,
         ),
     ]
 )
@@ -48,7 +48,7 @@ POS_SHIFTS = torch.tensor(
         [0, 1, 0, 0],
         [0, 0, 1, 0],
     ],
-    dtype=torch.int32,
+    dtype=torch.int64,
 )
 
 
@@ -66,30 +66,30 @@ FACE_ROTATIONS = torch.stack(
 
 def build_actions_tensor(size: int) -> torch.Tensor:
     """
-    Built the 5D tensor carrying all rotations of a cube as matrix multiplication.
+    Built the 4D tensor carrying all rotations of a cube as index permutation.
     """
-    return torch.stack(
+    return torch.tensor(
         [
-            build_action_tensor(size=size, axis=axis, slice=slice, inverse=inverse)
+            [
+                [build_action_permutation(size=size, axis=axis, slice=slice, inverse=inverse) for inverse in range(2)]
+                for slice in range(size)
+            ]
             for axis in range(3)
-            for slice in range(size)
-            for inverse in range(2)
         ],
-        dim=0,
-    ).sum(dim=0, dtype=torch.int32)
+        dtype=torch.int64,
+    )
 
 
-def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch.Tensor:
+def build_action_permutation(size: int, axis: int, slice: int, inverse: int) -> list[int]:
     """
-    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.
+    Compute the permutation list whose effect on a position-frozen color vector is the rotation
+    along the specified axis, within the specified slice and the specified orientation.
     """
-    tensor = build_cube_tensor(size).to(dtype=torch.int32)
+    tensor = build_cube_tensor(size).to(dtype=torch.int64)
     length = 6 * (size**2)
 
     # extract faces impacted by the move
-    indices = tensor.indices().to(dtype=torch.int32)  # size = (4, length)
+    indices = tensor.indices().to(dtype=torch.int64)  # size = (4, length)
     changes = (indices[axis + 1] == slice).nonzero().reshape(-1)  # size = (n,), n < length
     extract = indices[:, changes]  # size = (4, n)
 
@@ -99,7 +99,7 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     rotated = rotated + offsets  # size = (4, n)
 
     # apply face rotation
-    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(torch.int32) @ FACE_ROTATIONS[axis]).argmax(dim=-1)
+    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(torch.int64) @ 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)
@@ -107,28 +107,12 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     outputs = outputs.transpose(0, 1).tolist()  # size = (n, 4)
 
     # compute position-based permutation of colors equivalent to rotation converting inputs into outputs
+    local_perm = {i: outputs.index(inputs[i]) for i in range(len(inputs))}
     local_to_total = dict(enumerate(changes.tolist()))
     total_to_local = {ind: i for i, ind in local_to_total.items()}
 
-    local_perm = {i: outputs.index(inputs[i]) for i in range(len(inputs))}
-    total_perm = {
-        i: (i if i not in total_to_local else local_to_total[local_perm[total_to_local[i]]]) for i in range(length)
-    }
-
-    # 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=torch.int32,
-    )
-    perm_values = torch.tensor([1] * length, dtype=torch.int32)
-    perm_size = (3, size, 2, length, length)
-    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=torch.int32)
+    # return permutation on total list of facelet positions
+    return [(i if i not in total_to_local else local_to_total[local_perm[total_to_local[i]]]) for i in range(length)]
 
 
 def parse_action_str(move: str) -> tuple[int, int, int]:

src/rubik/cube.py

@@ -2,7 +2,6 @@ from functools import reduce
 from loguru import logger
 
 import torch
-import torch.nn.functional as F
 
 from rubik.action import build_actions_tensor, parse_actions_str, sample_actions_str
 from rubik.state import build_cube_tensor
@@ -26,10 +25,11 @@ class Cube:
         """
         tensor = build_cube_tensor(size)
 
-        self.dtype = torch.int8 if size <= 6 else torch.int16 if size <= 73 else torch.int32
-        self.coordinates = tensor.indices().to(self.dtype)
-        self.state = F.one_hot(tensor.values().long(), num_classes=7).to(self.dtype)
-        self.actions = build_actions_tensor(size).to(self.dtype)
+        self.dtype = torch.int64
+        self.coordinates = tensor.indices()
+        self.state = tensor.values()
+        self.actions = build_actions_tensor(size)
+        # internal-only attributes
         self._history: list[tuple[int, int, int]] = []
         self._colors: list[str] = list("ULCRBD")
         self._size: int = size
@@ -54,7 +54,7 @@ class Cube:
         """
         tensor = torch.sparse_coo_tensor(
             indices=self.coordinates,
-            values=self.state.argmax(dim=-1),
+            values=self.state,
             size=(6, self.size, self.size, self.size),
             dtype=self.dtype,
         ).to_dense()
@@ -72,13 +72,7 @@ class Cube:
 
     def to(self, device: str | torch.device) -> "Cube":
         device = torch.device(device)
-        dtype = (
-            self.state.dtype
-            if self.state.device == device
-            else self.dtype
-            if device == torch.device("cpu")
-            else torch.float32
-        )
+        dtype = self.dtype if device == torch.device("cpu") else torch.float32
         self.state = self.state.to(device=device, dtype=dtype)
         self.actions = self.actions.to(device=device, dtype=dtype)
         logger.info(f"Using device '{self.state.device}' and dtype '{dtype}'")
@@ -114,18 +108,17 @@ class Cube:
         Apply a move (defined as 3 coordinates) to the cube.
         """
         action = self.actions[axis, slice, inverse]
-        self.state = action @ self.state
+        self.state = torch.gather(self.state, 0, action)
         self._history.append((axis, slice, inverse))
         return
 
-    def compute_changes(self, moves: str) -> dict[int, int]:
+    def compose_moves(self, moves: str) -> torch.Tensor:
         """
         combine a sequence of moves and return the resulting changes.
         """
         actions = parse_actions_str(moves)
-        tensors = [self.actions[*action].to(torch.float32) for action in actions]
-        result = reduce(lambda A, B: B @ A, tensors).to(torch.int16).coalesce()
-        return dict(result.indices().transpose(0, 1).tolist())
+        tensors = [self.actions[*action] for action in actions]
+        return reduce(lambda A, B: torch.gather(A, 0, B), tensors)
 
     def __str__(self):
         """

src/rubik/interface/plot.py

@@ -7,6 +7,7 @@ import torch
 class CubeVisualizer:
     """
     Utility class for ploting a cube, with some layout ingredients precomputed at init.
+    Greatly inspired from https://www.kaggle.com/code/edomingo/nxn-rubik-s-cube-3d-interactive-viz-plotly/notebook.
     """
 
     def __init__(self, size: int):
@@ -108,7 +109,7 @@ class CubeVisualizer:
         Generates a 3D plot of a cube given its coordinates, state and size.
         """
         # set the color of each facelet, face after face
-        face_state = (state.argmax(dim=-1) - 1).reshape(6, -1).tolist()
+        face_state = (state - 1).reshape(6, -1).tolist()
         face_colors = [[self.colors[f] for f in face] for face in face_state]
 
         face_coordinates = [coordinates[1:, (coordinates[0] == i)].transpose(0, 1).tolist() for i in range(6)]

src/rubik/state.py

@@ -9,7 +9,7 @@ def build_cube_tensor(size: int) -> torch.Tensor:
 
     # build dense tensor filled with colors
     n = size - 1
-    tensor = torch.zeros([6, size, size, size], dtype=torch.int32)
+    tensor = torch.zeros([6, size, size, size], dtype=torch.int64)
     tensor[0, :, :, n] = 1  # up
     tensor[1, 0, :, :] = 2  # left
     tensor[2, :, n, :] = 3  # front
@@ -25,6 +25,6 @@ def build_permutation_matrix(size: int, perm: str) -> torch.Tensor:
     """
     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.int32)
-    values = torch.tensor([1] * size, dtype=torch.int32)
-    return torch.sparse_coo_tensor(indices=indices, values=values, size=(size, size), dtype=torch.int32).coalesce()
+    indices = torch.tensor([list(range(size)), [(perm_dict.get(i, i)) for i in range(size)]], dtype=torch.int64)
+    values = torch.tensor([1] * size, dtype=torch.int64)
+    return torch.sparse_coo_tensor(indices=indices, values=values, size=(size, size), dtype=torch.int64).coalesce()

tests/unit/test_action.py

@@ -8,7 +8,7 @@ from rubik.action import (
     POS_SHIFTS,
     FACE_ROTATIONS,
     build_actions_tensor,
-    build_action_tensor,
+    build_action_permutation,
     parse_action_str,
     parse_actions_str,
     sample_actions_str,
@@ -96,7 +96,7 @@ def test_build_actions_tensor_shape(size: int):
     """
     Test that "build_actions_tensor" output has expected shape.
     """
-    expected = (3, size, 2, 6 * (size**2), 6 * (size**2))
+    expected = (3, size, 2, 6 * (size**2))
     observed = build_actions_tensor(size).shape
     assert expected == observed, (
         f"'build_actions_tensor' output has incorrect shape: expected shape '{expected}', got '{observed}' instead"
@@ -111,14 +111,14 @@ def test_build_actions_tensor_shape(size: int):
         (5, 1, 4, 0),
     ],
 )
-def test_build_action_tensor_shape(size: int, axis: int, slice: int, inverse: int):
+def test_build_action_permutation(size: int, axis: int, slice: int, inverse: int):
     """
     Test that "build_actions_tensor" output has expected shape.
     """
-    expected = (3, size, 2, 6 * (size**2), 6 * (size**2))
-    observed = build_action_tensor(size, axis, slice, inverse).shape
+    expected = 6 * (size**2)
+    observed = len(build_action_permutation(size, axis, slice, inverse))
     assert expected == observed, (
-        f"'build_action_tensor' output has incorrect shape: expected shape '{expected}', got '{observed}' instead"
+        f"'build_action_tensor' output has incorrect length: expected length '{expected}', got '{observed}'"
     )
 
 

tests/unit/test_cube.py

@@ -16,8 +16,8 @@ class TestCube:
         Test that the __init__ method produce expected attributes.
         """
         cube = Cube(size)
-        assert cube.state.shape == (6 * (size**2), 7), f"'state' has incorrect shape {cube.state.shape}"
-        assert cube.actions.shape == (3, size, 2, cube.state.shape[0], cube.state.shape[0]), (
+        assert cube.state.shape == (6 * (size**2),), f"'state' has incorrect shape {cube.state.shape}"
+        assert cube.actions.shape == (3, size, 2, cube.state.shape[0]), (
             f"'actions' has incorrect shape {cube.actions.shape}"
         )
         assert len(cube.history) == 0, "'history' field should be empty"
@@ -93,23 +93,22 @@ class TestCube:
             "X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i " * 2,
         ],
     )
-    def test_compute_changes(self, moves: str):
+    def test_compose_moves(self, moves: str):
         """
-        Test that the .compute_changes method behaves as expected.
+        Test that the .compose_moves method behaves as expected.
         """
         cube = Cube(3)
-        facets = cube.state.argmax(dim=-1).to(cube.dtype).tolist()
-        changes = cube.compute_changes(moves)
 
         # apply changes induced by moves using the permutation dict returned by 'compute_changes'
-        expected = [facets[changes.get(i, i)] for i in range(len(facets))]
+        changes = cube.compose_moves(moves)
+        expected = torch.gather(cube.state.clone(), 0, changes)
 
         # apply changes induced by moves using the optimized 'rotate' method
         cube.rotate(moves)
-        observed = cube.state.argmax(dim=-1).to(cube.dtype).tolist()
+        observed = cube.state
 
         # assert the tow are identical
-        assert expected == observed, "method 'compute_changes' does not behave correctly: "
+        assert torch.equal(expected, observed), "method 'compute_changes' does not behave correctly: "
 
     def test__str__len(self):
         """