add builder for tensor of actions

README.md

@@ -17,7 +17,7 @@ pre-commit install
 ```python
 from rubik.cube import Cube
 
-cube = Cube.from_default(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
+cube = Cube.create(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
 print(cube)
 #     UUU        
 #     UUU

src/rubik/__main__.py

@@ -1,6 +1,4 @@
-from fire import Fire
+# from fire import Fire
 
-from rubik.hello import hello_world
 
-
-Fire({"hello": hello_world})
+# Fire({"hello": hello_world})

src/rubik/cube.py

@@ -6,38 +6,50 @@ import torch
 @dataclass
 class Cube:
     """
-    A 5D tensor filled with 0 or 1. Dimensions have the following interpretation:
+    A 4D tensor filled with colors. Dimensions have the following interpretation:
+        - Face (from 0 to 5, with 0 = "Up", 1 = "Left", 2 = "Front", 3 = "Right", 4 = "Back", 5 = "Down").
         - X coordinate (from 0 to self.size - 1, from Left to Right).
         - Y coordinate (from 0 to self.size - 1, from Back to Front).
         - Z coordinate (from 0 to self.size - 1, from Down to Up).
-        - Face (from 0 to 5, with 0 = "Up", 1 = "Left", 2 = "Front", 3 = "Right", 4 = "Back", 5 = "Down").
-        - Color (from 0 to 6, 0 being the "dark" color, the rest according to order given in "colors" attribute).
+
+    Colors filling each tensor cell are from 0 to 6, 0 being the "dark" color,
+    the rest according to order given in "colors" attribute.
     """
 
+    coordinates: torch.Tensor
     state: torch.Tensor
     colors: list[str]
     size: int
 
     @classmethod
-    def from_default(cls, colors: list[str], size: int) -> "Cube":
+    def create(cls, colors: list[str], size: int) -> "Cube":
         """
         Create Cube from a given list of 6 colors and size.
         Example:
-            cube = Cube.from_default(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
+            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 tensor filled with 0's, and fill the faces with 1's
+        # build dense tensor filled with colors
         n = size - 1
-        state = torch.zeros([size, size, size, 6, 7], dtype=torch.int8)
-        state[:, :, n, 0, 1] = 1  # up
-        state[0, :, :, 1, 2] = 1  # left
-        state[:, 0, :, 2, 3] = 1  # front
-        state[n, :, :, 3, 4] = 1  # right
-        state[:, n, :, 4, 5] = 1  # back
-        state[:, :, 0, 5, 6] = 1  # down
-        return cls(state, colors, size)
+        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
+
+    def rotate(self, moves: str):
+        raise NotImplementedError
+
+    def solve(slef, policy: str):
+        raise NotImplementedError
 
     @staticmethod
     def pad_colors(colors: list[str]) -> list[str]:
@@ -47,41 +59,35 @@ class Cube:
         max_len = max(len(c) for c in colors)
         return [c + " " * (max_len - len(c)) for c in colors]
 
-    def to_grid(self, pad_colors: bool = False) -> list[list[list[str]]]:
+    @classmethod
+    def from_sparse(cls, tensor: torch.Tensor, colors: list[str], size: int) -> "Cube":
         """
-        Convert Cube into a 3D grid representation.
+        Gather cube attributes into a torch sparse tensor.
         """
-        n = self.size - 1
-        colors = self.pad_colors(self.colors) if pad_colors else self.colors
-        grid = [
-            self.state[:, :, n, 0, :].argmax(dim=-1),  # up
-            self.state[0, :, :, 1, :].argmax(dim=-1),  # left
-            self.state[:, 0, :, 2, :].argmax(dim=-1),  # front
-            self.state[n, :, :, 3, :].argmax(dim=-1),  # right
-            self.state[:, n, :, 4, :].argmax(dim=-1),  # back
-            self.state[:, :, 0, 5, :].argmax(dim=-1),  # down
-        ]
-        return [[[colors[i - 1] for i in row] for row in face.tolist()] for face in grid]
+        coordinates = tensor.indices().transpose(0, 1).to(torch.int8)
+        values = tensor.values()
+        return cls(coordinates, values, colors, size)
+
+    def to_sparse(self) -> torch.Tensor:
+        """
+        Gather cube attributes into a torch sparse tensor.
+        """
+        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,
+        )
 
     def __str__(self):
         """
         Compute a string representation of a cube.
-        Example:
-            cube = Cube.from_default(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
-            print(cube)
-            #     UUU
-            #     UUU
-            #     UUU
-            # LLL CCC RRR BBB
-            # LLL CCC RRR BBB
-            # LLL CCC RRR BBB
-            #     DDD
-            #     DDD
-            #     DDD
         """
-        grid = self.to_grid(pad_colors=True)
+        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 grid[0])
-        l2 = "\n".join(" ".join("".join(grid[face_i][row_i]) for face_i in range(1, 5)) for row_i in range(self.size))
-        l3 = "\n".join(" ".join((void, "".join(row), void, void)) for row in grid[-1])
+        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])

src/rubik/hello.py

@@ -1,2 +0,0 @@
-def hello_world(s: str) -> None:
-    print(f"Hello {s} !")

src/rubik/moves.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn.functional as F
+
+from rubik.cube import Cube
+
+
+INT8 = torch.int8
+
+POS_ROTATIONS = torch.stack(
+    [
+        # rot about X: Z -> Y
+        torch.tensor(
+            [
+                [1, 0, 0, 0],
+                [0, 1, 0, 0],
+                [0, 0, 0, 1],
+                [0, 0, -1, 0],
+            ],
+            dtype=INT8,
+        ),
+        # rot about Y: X -> Z
+        torch.tensor(
+            [
+                [1, 0, 0, 0],
+                [0, 0, 0, -1],
+                [0, 0, 1, 0],
+                [0, 1, 0, 0],
+            ],
+            dtype=INT8,
+        ),
+        # rot about Z: Y -> X
+        torch.tensor(
+            [
+                [1, 0, 0, 0],
+                [0, 0, 1, 0],
+                [0, -1, 0, 0],
+                [0, 0, 0, 1],
+            ],
+            dtype=INT8,
+        ),
+    ]
+)
+
+POS_SHIFTS = torch.tensor(
+    [
+        [0, 0, 0, 2],
+        [0, 2, 0, 0],
+        [0, 0, 2, 0],
+    ],
+    dtype=INT8,
+)
+
+FACE_PERMS = 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,
+        ),
+    ]
+).transpose(1, 2)
+
+
+def build_actions_tensor(size: int) -> torch.Tensor:
+    """
+    Built the 5D tensor carrying all rotations of a cube as matrix multiplication.
+    """
+    return torch.stack(
+        [
+            build_permunation_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)
+
+
+def build_permunation_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)
+    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)
+
+    # apply coordinate rotation
+    rotated = POS_ROTATIONS[axis] @ extract  # size = (4, n)
+    offsets = POS_SHIFTS[axis].repeat(extract.shape[-1], 1).transpose(0, 1)  # size = (4, n)
+    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)
+
+    # from this point on, convert rotation into a position-based permutation of colors
+    (inputs, outputs) = (rotated, extract) if bool(inverse) else (extract, rotated)
+    inputs = inputs.transpose(0, 1).tolist()  # size = (n, 4)
+    outputs = outputs.transpose(0, 1).tolist()  # size = (n, 4)
+
+    extract_to_coordinates = dict(enumerate(indices.tolist()))
+    coordinates_to_extract = {ind: i for i, ind in extract_to_coordinates.items()}
+
+    extract_perm = {i: inputs.index(outputs[i]) for i in range(len(inputs))}
+    global_perm = {
+        i: (i if i not in coordinates_to_extract else extract_to_coordinates[extract_perm[coordinates_to_extract[i]]])
+        for i in range(length)
+    }
+    perm_indices = torch.tensor(
+        [[axis] * length, [slice] * length, [inverse] * length, list(global_perm.keys()), list(global_perm.values())],
+        dtype=INT8,
+    )
+    perm_values = torch.tensor([1] * length)
+    perm_size = (3, size, 2, length, length)
+    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=INT8)