Gradio Interface (#2)

* add plotly and gradio

* edit readme

* simplify cube object

* simplify cube object

* edit facelect computation

* base demo with 3d plot

* remove legacy plot function

* rename demo as interface

* edit plot function

* remove cupy dependency

* all dependencies

* gradio interface with 3D plot

README.md

@@ -14,14 +14,20 @@ pre-commit install
 
 ## Usage
 
-### Create a cube
+### Launch the web interface
+
+```shell
+python -m rubik interface
+```
+
+### Use the python API
 
 ```python
 from rubik.cube import Cube
 
-cube = Cube(colors=['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
+cube = Cube(size=3)
 
-# display the cube state and history of moves
+# display cube state
 print(cube)
 #     UUU        
 #     UUU
@@ -33,17 +39,14 @@ print(cube)
 #     DDD
 #     DDD
 
+# display history of moves
 print(cube.history)
 # []
-```
-
-### Perform basic moves
 
-```python
-# shuffle the cube using 1000 random moves (random shuffling resets the history)
-cube.shuffle(num_moves=1000, seed=0)
+# scramble the cube using 1000 random moves (this resets the history)
+cube.scramble(num_moves=1000, seed=0)
 
-# rotate it in some way
+# rotate it in some way (this gets appended to history)
 cube.rotate('X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i')
 ```
 
@@ -66,8 +69,9 @@ cube.rotate('X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i')
 
 
 
+## References
 
-## Related projects
+### Python implementations & rule-based solvers
 
 Open-source projects related to Rubik's Cube, sorted by number of stars:
 - [adrianliaw/PyCuber](https://github.com/adrianliaw/PyCuber)
@@ -77,3 +81,16 @@ 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)
+
+### Machine Learning based solver models
+
+- 2025, _CayleyPy Cube_, [Github](https://github.com/k1242/cayleypy-cube), [Paper](https://arxiv.org/html/2502.13266v1)
+
+- 2025, _Solving A Rubik’s Cube with Supervised Learning – Intuitively and Exhaustively Explained_, [Blog post](https://towardsdatascience.com/solving-a-rubiks-cube-with-supervised-learning-intuitively-and-exhaustively-explained-4f87b72ba1e2/)
+
+- 2024, _Solving Rubik's Cube Without Tricky Sampling_, [Paper](https://arxiv.org/abs/2411.19583).<br>
+This involves training a scorer, that estimates the number of moves transforming a given source state into a given target state, where the latter is not necessarily a solved cube. Data are synthetically generated performing random moves and factorizing repeated identical moves.
+
+- 2023, _Curious Transformer_, [Github](https://github.com/tedtedtedtedtedted/Solve-Rubiks-Cube-Via-Transformer)
+
+- 2021, _Efficient Cube_, [Github](https://github.com/kyo-takano/efficientcube), [Paper](https://arxiv.org/abs/2106.03157)

notebooks/dev.ipynb

@@ -31,7 +31,8 @@
    "outputs": [],
    "source": [
     "from rubik.cube import Cube\n",
-    "from rubik.action import build_actions_tensor"
+    "from rubik.action import build_actions_tensor\n",
+    "from rubik.interface.plot import CubeVisualizer"
    ]
   },
   {
@@ -45,7 +46,7 @@
     "\n",
     "actions = build_actions_tensor(size)\n",
     "\n",
-    "cube = Cube([\"U\", \"L\", \"C\", \"R\", \"B\", \"D\"], size=size)\n",
+    "cube = Cube(size)\n",
     "print(cube)"
    ]
   },
@@ -57,7 +58,8 @@
    "outputs": [],
    "source": [
     "cubis = copy.deepcopy(cube)\n",
-    "cubis.shuffle(2000, seed=0)\n",
+    "cubis.rotate(\"X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i \")\n",
+    "\n",
     "print(cubis)\n",
     "print(cubis.history)"
    ]
@@ -69,7 +71,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cubis.compute_changes(\"X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i \")"
+    "visualizer = CubeVisualizer(size=cubis.size)\n",
+    "layout_args = {\"autosize\": False, \"width\": 600, \"height\": 600}\n",
+    "\n",
+    "visualizer(cubis.coordinates, cubis.state, cubis.size).update_layout(**layout_args).show()"
    ]
   },
   {
@@ -80,7 +85,7 @@
    "outputs": [],
    "source": [
     "cubis = copy.deepcopy(cube)\n",
-    "cubis.rotate(\"X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i \" * 1000)\n",
+    "cubis.scramble(2000, seed=0)\n",
     "print(cubis)\n",
     "print(cubis.history)"
    ]
@@ -91,6 +96,19 @@
    "id": "7",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "cubis = copy.deepcopy(cube)\n",
+    "cubis.rotate(\"X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i \" * 1000)\n",
+    "print(cubis)\n",
+    "print(cubis.history)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "(actions[0, 2, 0].type(torch.float32) @ actions[0, 1, 1].type(torch.float32)).type(torch.int8)"
    ]

pyproject.toml

@@ -5,9 +5,11 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.11,<3.12"
 dependencies = [
-    "cupy-cuda12x>=13.4.1",
     "fire>=0.7.0",
+    "gradio>=5.38",
     "loguru>=0.7.3",
+    "plotly>=6.2.0",
+    "pydantic>=2.11.7",
     "torch>=2.7.1",
 ]
 

src/rubik/__main__.py

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

src/rubik/action.py

@@ -4,7 +4,7 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 
-from rubik.tensor_utils import build_permutation_matrix, build_cube_tensor
+from rubik.state import build_permutation_matrix, build_cube_tensor
 
 
 POS_ROTATIONS = torch.stack(
@@ -17,7 +17,7 @@ POS_ROTATIONS = torch.stack(
                 [0, 0, 0, 1],
                 [0, 0, -1, 0],
             ],
-            dtype=torch.int16,
+            dtype=torch.int32,
         ),
         # 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.int16,
+            dtype=torch.int32,
         ),
         # 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.int16,
+            dtype=torch.int32,
         ),
     ]
 )
@@ -48,7 +48,7 @@ POS_SHIFTS = torch.tensor(
         [0, 1, 0, 0],
         [0, 0, 1, 0],
     ],
-    dtype=torch.int16,
+    dtype=torch.int32,
 )
 
 
@@ -76,7 +76,7 @@ def build_actions_tensor(size: int) -> torch.Tensor:
             for inverse in range(2)
         ],
         dim=0,
-    ).sum(dim=0, dtype=torch.int16)
+    ).sum(dim=0, dtype=torch.int32)
 
 
 def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch.Tensor:
@@ -85,11 +85,11 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     is the rotation along the specified axis, within the specified slice and the specified
     orientation.
     """
-    tensor = build_cube_tensor(colors=list("ULCRBD"), size=size)
+    tensor = build_cube_tensor(size).to(dtype=torch.int32)
     length = 6 * (size**2)
 
     # extract faces impacted by the move
-    indices = tensor.indices().to(dtype=torch.int16)  # size = (4, length)
+    indices = tensor.indices().to(dtype=torch.int32)  # 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.int16) @ FACE_ROTATIONS[axis]).argmax(dim=-1)
+    rotated[0] = (F.one_hot(rotated[0].long(), num_classes=6).to(torch.int32) @ 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)
@@ -110,19 +110,25 @@ def build_action_tensor(size: int, axis: int, slice: int, inverse: int) -> torch
     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))}
+    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.int16,
+        [
+            [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.int16)
+    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.int16)
+    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=torch.int32)
 
 
 def parse_action_str(move: str) -> tuple[int, int, int]:

src/rubik/cube.py

@@ -5,8 +5,7 @@ import torch
 import torch.nn.functional as F
 
 from rubik.action import build_actions_tensor, parse_actions_str, sample_actions_str
-from rubik.display import stringify
-from rubik.tensor_utils import build_cube_tensor
+from rubik.state import build_cube_tensor
 
 
 class Cube:
@@ -21,30 +20,65 @@ class Cube:
     the rest according to order given in "colors" attribute.
     """
 
-    def __init__(self, colors: list[str], size: int):
+    def __init__(self, size: int):
         """
-        Create Cube from a given list of 6 colors and size.
-        Example:
-            cube = Cube(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
+        Create Cube of a given size.
         """
-        tensor = build_cube_tensor(colors, size)
-        self.coordinates = tensor.indices().transpose(0, 1).to(torch.int16)
-        self.state = F.one_hot(tensor.values().long(), num_classes=7).to(torch.int16)
-        self.actions = build_actions_tensor(size)
-        self.history: list[list[int]] = []
-        self.colors = colors
-        self.size = size
+        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._history: list[tuple[int, int, int]] = []
+        self._colors: list[str] = list("ULCRBD")
+        self._size: int = size
+
+    @property
+    def history(self) -> list[tuple[int, int, int]]:
+        return self._history
+
+    @property
+    def colors(self) -> list[str]:
+        return self._colors
+
+    @property
+    def size(self) -> int:
+        return self._size
+
+    @property
+    def facelets(self) -> list[list[list[str]]]:
+        """
+        Return the list of faces of the cube, each given by a list of rows,
+        each given by a list of facelets.
+        """
+        tensor = torch.sparse_coo_tensor(
+            indices=self.coordinates,
+            values=self.state.argmax(dim=-1),
+            size=(6, self.size, self.size, self.size),
+            dtype=self.dtype,
+        ).to_dense()
+
+        n = self.size - 1
+        faces = [
+            tensor[0, :, :, n].transpose(0, 1),  # up
+            tensor[1, 0, :, :].flip(1).transpose(0, 1),  # left
+            tensor[2, :, n, :].flip(1).transpose(0, 1),  # front
+            tensor[3, n, :, :].flip(0).flip(1).transpose(0, 1),  # right
+            tensor[4, :, 0, :].flip(0).flip(1).transpose(0, 1),  # back
+            tensor[5, :, :, 0].flip(1).transpose(0, 1),  # down
+        ]
+        return [[[self.colors[i - 1] for i in row] for row in face.tolist()] for face in faces]
 
     def to(self, device: str | torch.device) -> "Cube":
         device = torch.device(device)
         dtype = (
             self.state.dtype
             if self.state.device == device
-            else torch.int16
+            else self.dtype
             if device == torch.device("cpu")
             else torch.float32
         )
-        self.coordinates = self.coordinates.to(device=device, dtype=dtype)
         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}'")
@@ -54,10 +88,10 @@ class Cube:
         """
         Reset internal history of moves.
         """
-        self.history = []
+        self._history = []
         return
 
-    def shuffle(self, num_moves: int, seed: int = 0) -> None:
+    def scramble(self, num_moves: int, seed: int = 0) -> None:
         """
         Randomly shuffle the cube by the supplied number of steps, and reset history of moves.
         """
@@ -81,7 +115,7 @@ class Cube:
         """
         action = self.actions[axis, slice, inverse]
         self.state = action @ self.state
-        self.history.append([axis, slice, inverse])
+        self._history.append((axis, slice, inverse))
         return
 
     def compute_changes(self, moves: str) -> dict[int, int]:
@@ -97,5 +131,9 @@ class Cube:
         """
         Compute a string representation of a cube.
         """
-        state = self.state.argmax(dim=-1).to(device="cpu", dtype=torch.int16)
-        return stringify(state, self.colors, self.size)
+        space = " " * self.size
+        facelets = self.facelets
+        l1 = "\n".join(" ".join([space, "".join(row), space, space]) for row in facelets[0])
+        l2 = "\n".join(" ".join("".join(face[i]) for face in facelets[1:5]) for i in range(self.size))
+        l3 = "\n".join(" ".join((space, "".join(row), space, space)) for row in facelets[-1])
+        return "\n".join([l1, l2, l3])

src/rubik/display.py

@@ -1,23 +0,0 @@
-import torch
-
-
-def stringify(state: torch.Tensor, colors: list[str], size: int) -> str:
-    """
-    Compute a string representation of a cube.
-    """
-    colors = pad_colors(colors)
-    faces = state.reshape(6, size, 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 colors) * 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(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/interface/app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+
+from plotly import graph_objects as go
+
+from rubik.cube import Cube
+from rubik.interface.plot import CubeVisualizer
+
+
+def app(default_size: int = 3):
+    """
+    Interface with the following features:
+        - create a cube of the specified size.
+        - ability to scramble it with a specified number of moves.
+        - ability to rotate it through a text field.
+        - display a cube upon creation or update.
+    """
+    cube = Cube(default_size)
+    cube_visualizer = CubeVisualizer(default_size)
+
+    def create_cube(size) -> None:
+        nonlocal cube
+        nonlocal cube_visualizer
+        cube = Cube(size)
+        cube_visualizer = CubeVisualizer(size)
+        return
+
+    def scramble_cube(num_moves: int) -> None:
+        nonlocal cube
+        cube.scramble(num_moves, seed=0)
+        return
+
+    def rotate_cube(moves: str) -> None:
+        nonlocal cube
+        cube.rotate(moves)
+        return
+
+    def display_cube() -> go.Figure:
+        nonlocal cube
+        nonlocal cube_visualizer
+        layout_args = {"autosize": False, "width": 600, "height": 600}
+        return cube_visualizer(cube.coordinates, cube.state, cube.size).update_layout(**layout_args)
+
+    with gr.Blocks(fill_height=True) as demo:
+        # structure
+        gr.Markdown("Rubik's Cube Interface")
+        with gr.Row():
+            with gr.Column(scale=15):
+                size = gr.Slider(1, 100, value=default_size, step=1, label="Select a size")
+                create_btn = gr.Button("Generate a Cube")
+
+                num_moves = gr.Slider(0, 10000, value=500, step=100, label="Select a number of steps for scrambling")
+                scramble_btn = gr.Button("Scramble the Cube")
+
+                moves = gr.Textbox(value="X0 Y1 Z0i", label="Define a sequence of moves")
+                rotate_btn = gr.Button("Rotate the Cube")
+
+            with gr.Column(scale=85):
+                plot = gr.Plot(display_cube(), container=False)
+
+        # interactions
+        create_btn.click(create_cube, inputs=size).success(display_cube, None, plot)
+        scramble_btn.click(scramble_cube, inputs=num_moves).success(display_cube, None, plot)
+        rotate_btn.click(rotate_cube, inputs=moves).success(display_cube, None, plot)
+
+    demo.launch()
+    return

src/rubik/interface/plot.py

@@ -0,0 +1,155 @@
+import copy
+import plotly.graph_objects as go
+
+import torch
+
+
+class CubeVisualizer:
+    """
+    Utility class for ploting a cube, with some layout ingredients precomputed at init.
+    """
+
+    def __init__(self, size: int):
+        self.vertices = self.build_vertices(size)
+        self.shifts = self.build_shifts()
+        self.x_coor = [v[1] for v in self.vertices]
+        self.y_coor = [v[0] for v in self.vertices]
+        self.z_coor = [v[2] for v in self.vertices]
+        self.fig = self.build_base_figure(size)
+
+    @property
+    def colors(self):
+        """
+        Assign colors for each face.
+        """
+        return {
+            0: "white",
+            1: "#3588cc",
+            2: "red",
+            3: "green",
+            4: "orange",
+            5: "yellow",
+        }
+
+    @staticmethod
+    def build_vertices(size: int):
+        face_vertices = [
+            [[x, y, size] for y in range(size + 1) for x in range(size + 1)],  # Up
+            [[0, y, z] for y in range(size + 1) for z in range(size + 1)],  # Left
+            [[x, size, z] for x in range(size + 1) for z in range(size + 1)],  # Front
+            [[size, y, z] for y in range(size + 1) for z in range(size + 1)],  # Right
+            [[x, 0, z] for x in range(size + 1) for z in range(size + 1)],  # Back
+            [[x, y, 0] for x in range(size + 1) for y in range(size + 1)],  # Down
+        ]
+        return [vertex for face in face_vertices for vertex in face]
+
+    @staticmethod
+    def build_shifts():
+        return [
+            [(0, 1, 0), (1, 0, 0), (1, 1, 0)],  # Up
+            [(0, 1, 0), (0, 0, 1), (0, 1, 1)],  # Left
+            [(1, 0, 0), (0, 0, 1), (1, 0, 1)],  # Front
+            [(0, 1, 0), (0, 0, 1), (0, 1, 1)],  # Right
+            [(1, 0, 0), (0, 0, 1), (1, 0, 1)],  # Back
+            [(0, 1, 0), (1, 0, 0), (1, 1, 0)],  # Down
+        ]
+
+    @staticmethod
+    def build_base_figure(size):
+        """
+        Create base figure for the cube, containing everything but the cube facelets.
+        """
+        fig = go.Figure()
+
+        # add black lines to the cube
+        lines_seq = [[0, size, size, 0, 0], [0, 0, size, size, 0]]
+        lines_args = {"mode": "lines", "line": {"width": 5, "color": "black"}, "hoverinfo": "none"}
+        for i in range(size + 1):
+            fig.add_trace(go.Scatter3d(x=[i] * 5, y=lines_seq[1], z=lines_seq[0], **lines_args))
+            fig.add_trace(go.Scatter3d(x=lines_seq[1], y=[i] * 5, z=lines_seq[0], **lines_args))
+            fig.add_trace(go.Scatter3d(x=lines_seq[0], y=lines_seq[1], z=[i] * 5, **lines_args))
+
+        # add text along each axis
+        fig.add_trace(
+            go.Scatter3d(
+                x=[size / 2, size / 2, size + 1.5 + size * 0.5],
+                y=[size / 2, -1.5 - size * 0.5, size / 2],
+                z=[size + 1 + size * 0.5, size / 2, size / 2],
+                mode="text",
+                text=["UP", "LEFT", "FRONT"],
+                textposition="middle center",
+                textfont={"size": 20},
+                hoverinfo="none",
+            )
+        )
+
+        # remove legend, background, grid, ticks, etc.
+        scene_axis = {
+            "showgrid": False,
+            "zeroline": False,
+            "showticklabels": False,
+            "showbackground": False,
+            "title_text": "",
+            "showspikes": False,
+        }
+        return fig.update_layout(
+            showlegend=False,
+            autosize=True,
+            scene={
+                "xaxis": scene_axis,
+                "yaxis": scene_axis,
+                "zaxis": scene_axis,
+                "camera": {"eye": {"x": 0.8, "y": -1.2, "z": 0.7}},
+            },
+        )
+
+    def __call__(self, coordinates: torch.Tensor, state: torch.Tensor, size: int) -> go.Figure:
+        """
+        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_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)]
+
+        # for each facelet of a face, draw 2 complementary triangles covering it
+        i_coor = []
+        j_coor = []
+        k_coor = []
+        facecolor = []
+        for i in range(6):
+            face = face_coordinates[i]
+            if i == 0:
+                face = [[x, y, size] for x, y, z in face]
+            if i == 2:
+                face = [[x, size, z] for x, y, z in face]
+            if i == 3:
+                face = [[size, y, z] for x, y, z in face]
+
+            # add first triangle of each facelet
+            i_coor += [self.vertices.index(p) for p in face]
+            j_coor += [self.vertices.index([c + s for c, s in zip(p, self.shifts[i][0])]) for p in face]
+            k_coor += [self.vertices.index([c + s for c, s in zip(p, self.shifts[i][1])]) for p in face]
+
+            # add second triangle of each facelet
+            i_coor += [self.vertices.index([c + s for c, s in zip(p, self.shifts[i][2])]) for p in face]
+            j_coor += [self.vertices.index([c + s for c, s in zip(p, self.shifts[i][0])]) for p in face]
+            k_coor += [self.vertices.index([c + s for c, s in zip(p, self.shifts[i][1])]) for p in face]
+
+            facecolor += face_colors[i] * 2
+
+        fig = copy.deepcopy(self.fig)
+        return fig.add_trace(
+            go.Mesh3d(
+                x=self.x_coor,
+                y=self.y_coor,
+                z=self.z_coor,
+                i=i_coor,
+                j=j_coor,
+                k=k_coor,
+                facecolor=facecolor,
+                opacity=1,
+                hoverinfo="none",
+            )
+        )

src/rubik/{tensor_utils.py → state.py}

@@ -1,16 +1,15 @@
 import torch
 
 
-def build_cube_tensor(colors: list[str], size: int) -> torch.Tensor:
+def build_cube_tensor(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.int16)
+    tensor = torch.zeros([6, size, size, size], dtype=torch.int32)
     tensor[0, :, :, n] = 1  # up
     tensor[1, 0, :, :] = 2  # left
     tensor[2, :, n, :] = 3  # front
@@ -26,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.int16)
-    values = torch.tensor([1] * size, dtype=torch.int16)
-    return torch.sparse_coo_tensor(indices=indices, values=values, size=(size, size), dtype=torch.int16).coalesce()
+    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()

tests/unit/test_cube.py

@@ -10,23 +10,16 @@ class TestCube:
     A testing class for the Cube class.
     """
 
-    @pytest.mark.parametrize(
-        "colors, size",
-        [
-            [["U", "L", "C", "R", "B", "D"], 3],
-            [["Up", "Left", "Center", "Right", "Back", "Down"], 5],
-            [["A", "BB", "CCC", "DDDD", "EEEEE", "FFFFFF"], 10],
-        ],
-    )
-    def test__init__(self, colors: list[str], size: int):
+    @pytest.mark.parametrize("size", [3, 5, 10, 25])
+    def test__init__(self, size: int):
         """
         Test that the __init__ method produce expected attributes.
         """
-        cube = Cube(colors, size)
-        assert cube.coordinates.shape == (6 * (size**2), 4), (
-            f"'coordinates' has incorrect shape {cube.coordinates.shape}"
-        )
+        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]), (
+            f"'actions' has incorrect shape {cube.actions.shape}"
+        )
         assert len(cube.history) == 0, "'history' field should be empty"
 
     @pytest.mark.parametrize("device", ["cpu"])
@@ -34,7 +27,7 @@ class TestCube:
         """
         Test that the .to method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube = Cube(3)
         cube_2 = cube.to(device)
         assert torch.equal(cube.state, cube_2.state), "cube has different state after calling 'to' method"
 
@@ -42,7 +35,7 @@ class TestCube:
         """
         Test that the .reset_history method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube = Cube(3)
         cube.rotate("X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i")
         cube.reset_history()
         assert cube.history == [], "method 'reset_history' does not flush content"
@@ -52,9 +45,9 @@ class TestCube:
         """
         Test that the .shuffle method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube = Cube(3)
         cube_state = cube.state.clone()
-        cube.shuffle(num_moves, seed)
+        cube.scramble(num_moves, seed)
         assert cube.history == [], "method 'shuffle' does not flush content"
         assert not torch.equal(cube_state, cube.state), "method 'shuffle' does not change state"
 
@@ -69,7 +62,7 @@ class TestCube:
         """
         Test that the .rotate method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube = Cube(3)
         cube_state = cube.state.clone()
         cube.rotate(moves)
         assert cube.history != [], "method 'rotate' does not update history"
@@ -87,10 +80,10 @@ class TestCube:
         """
         Test that the .rotate_once method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube = Cube(3)
         cube_state = cube.state.clone()
         cube.rotate_once(axis, slice, inverse)
-        assert cube.history == [[axis, slice, inverse]], "method 'rotate_once' does not update history"
+        assert cube.history == [(axis, slice, inverse)], "method 'rotate_once' does not update history"
         assert not torch.equal(cube_state, cube.state), "method 'rotate_once' does not change state"
 
     @pytest.mark.parametrize(
@@ -104,8 +97,8 @@ class TestCube:
         """
         Test that the .compute_changes method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
-        facets = cube.state.argmax(dim=-1).to(torch.int16).tolist()
+        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'
@@ -113,15 +106,25 @@ class TestCube:
 
         # apply changes induced by moves using the optimized 'rotate' method
         cube.rotate(moves)
-        observed = cube.state.argmax(dim=-1).to(torch.int16).tolist()
+        observed = cube.state.argmax(dim=-1).to(cube.dtype).tolist()
 
         # assert the tow are identical
         assert expected == observed, "method 'compute_changes' does not behave correctly: "
 
-    def test__str__(self):
+    def test__str__len(self):
         """
         Test that the __str__ method behaves as expected.
         """
-        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube = Cube(3)
         repr = str(cube)
         assert len(repr), "__str__ method returns an empty representation"
+
+    @pytest.mark.parametrize("size", [3, 5, 8, 10])
+    def test__str__content(self, size: int):
+        """
+        Test that stringify behaves as expected.
+        """
+        cube = Cube(size=size)
+        repr = str(cube)
+        lens = {len(line) for line in repr.split("\n")}
+        assert len(lens) == 1, f"'stringify' lines have variable length: {lens}"

tests/unit/test_display.py

@@ -1,42 +0,0 @@
-import pytest
-
-import torch
-
-from rubik.cube import Cube
-from rubik.display import stringify, pad_colors
-
-
-@pytest.mark.parametrize(
-    "colors, size",
-    [
-        [["U", "L", "C", "R", "B", "D"], 3],
-        [["Up", "Left", "Center", "Right", "Back", "Down"], 5],
-        [["A", "BB", "CCC", "DDDD", "EEEEE", "FFFFFF"], 10],
-    ],
-)
-def test_stringify(colors: list[str], size: int):
-    """
-    Test that stringify behaves as expected.
-    """
-    cube = Cube(colors=colors, size=size)
-    state = cube.state.argmax(dim=-1).to(device="cpu", dtype=torch.int16)
-    repr = stringify(state, colors, size)
-    lens = {len(line) for line in repr.split("\n")}
-    assert len(lens) == 1, f"'stringify' lines have variable length: {lens}"
-
-
-@pytest.mark.parametrize(
-    "colors",
-    [
-        ["U", "L", "C", "R", "B", "D"],
-        ["Up", "Left", "Center", "Right", "Back", "Down"],
-        ["A", "BB", "CCC", "DDDD", "EEEEE", "FFFFFF"],
-    ],
-)
-def test_pad_colors(colors: list[str]):
-    """
-    Test that pad_colors behaves as expected.
-    """
-    padded = pad_colors(colors)
-    lengths = {len(color) for color in padded}
-    assert len(lengths) == 1, f"'pad_colors' generates non-unique lengths: {lengths}"

tests/unit/{test_tensor_utils.py → test_state.py}

@@ -2,7 +2,7 @@ import pytest
 
 import torch
 
-from rubik.tensor_utils import build_cube_tensor, build_permutation_matrix
+from rubik.state import build_cube_tensor, build_permutation_matrix
 
 
 @pytest.mark.parametrize("size", [2, 3, 5, 20])
@@ -10,7 +10,7 @@ def test_build_cube_tensor(size: int):
     """
     Test that build_cube_tensor behaves as expected.
     """
-    tensor = build_cube_tensor(colors=["U", "L", "C", "R", "B", "D"], size=size)
+    tensor = build_cube_tensor(size)
     facets = tensor.to_dense().to(dtype=torch.int8) != 0
     x_sums = facets.sum(dim=(0, 2, 3)).tolist()
     y_sums = facets.sum(dim=(0, 1, 3)).tolist()