Base cube and moves with sparse tensor multiplications

.gitignore

@@ -1,194 +1,12 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
 .coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# UV
-#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#uv.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
-.pdm.toml
-.pdm-python
-.pdm-build/
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
 .env
 .venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
-# Abstra
-# Abstra is an AI-powered process automation framework.
-# Ignore directories containing user credentials, local state, and settings.
-# Learn more at https://abstra.io/docs
-.abstra/
-
-# Visual Studio Code
-#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
-#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
-#  you could uncomment the following to ignore the enitre vscode folder
-# .vscode/
-
-# Ruff stuff:
-.ruff_cache/
+.mypy_cache
+.vscode
+.ruff_cache
+.ipynb_checkpoints
+.pytest_cache
 
-# PyPI configuration file
-.pypirc
+*egg-info
 
-# Cursor
-#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
-#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
-#  refer to https://docs.cursor.com/context/ignore-files
-.cursorignore
-.cursorindexingignore
+__pycache__

.pre-commit-config.yaml

@@ -0,0 +1,30 @@
+repos: 
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v5.0.0
+  hooks:
+    - id: check-added-large-files
+      args: [--maxkb=5000]
+    - id: detect-private-key
+
+- repo: https://github.com/astral-sh/uv-pre-commit
+  rev: 0.7.13
+  hooks:
+    - id: uv-lock
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.12.0
+  hooks:
+    - id: ruff
+      args: [--fix]
+    - id: ruff-format
+
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v1.16.1
+  hooks:
+    - id: mypy
+      exclude: '^(?!src).*'
+
+- repo: https://github.com/kynan/nbstripout
+  rev: 0.8.1
+  hooks:
+    - id: nbstripout

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -1 +1,79 @@
-# Rubik-Tensor
+# Rubik-Tensor
+
+
+## Setup
+
+This project uses `uv 0.7` as environment & dependency manager, and `python 3.11` as core interpreter. Install the project with
+
+```shell
+uv venv
+-- Activate env --
+uv sync
+pre-commit install
+```
+
+## Usage
+
+### Create a cube
+
+```python
+from rubik.cube import Cube
+
+cube = Cube(colors=['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
+
+# display the cube state and history of moves
+print(cube)
+#     UUU        
+#     UUU
+#     UUU
+# LLL CCC RRR BBB
+# LLL CCC RRR BBB
+# LLL CCC RRR BBB
+#     DDD
+#     DDD
+#     DDD
+
+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)
+
+# rotate it in some way
+cube.rotate('X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i')
+```
+
+## Roadmap
+
+#### Fully tensorized Rubik Cube model
+
+- ☑️ Tensorized states.
+- ☑️ Tensorized actions.
+- ☑️ Interface for performing actions.
+
+#### Movement explorer
+
+- ⬜ Explore changes resulting from a sequences of moves.
+- ⬜ Find least sequences of moves satisfying some input constrains.
+
+#### Visualization interface
+
+#### Base solvers following rule-based policies
+
+
+
+
+## Related projects
+
+Open-source projects related to Rubik's Cube, sorted by number of stars:
+- [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)
+- [bellerb/RubiksCube_Solver](https://github.com/bellerb/RubiksCube_Solver)
+- [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)

notebooks/dev.ipynb

@@ -0,0 +1,120 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy\n",
+    "\n",
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from rubik.cube import Cube\n",
+    "from rubik.action import build_actions_tensor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "size = 3\n",
+    "\n",
+    "actions = build_actions_tensor(size)\n",
+    "\n",
+    "cube = Cube([\"U\", \"L\", \"C\", \"R\", \"B\", \"D\"], size=size)\n",
+    "print(cube)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cubis = copy.deepcopy(cube)\n",
+    "cubis.shuffle(2000, seed=0)\n",
+    "print(cubis)\n",
+    "print(cubis.history)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cubis.compute_changes(\"X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i \")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "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": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "(actions[0, 2, 0].type(torch.float32) @ actions[0, 1, 1].type(torch.float32)).type(torch.int8)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Rubik-Tensor",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "rubik-tensor"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11,<3.12"
+dependencies = [
+    "cupy-cuda12x>=13.4.1",
+    "fire>=0.7.0",
+    "loguru>=0.7.3",
+    "torch>=2.7.1",
+]
+
+[dependency-groups]
+dev = [
+    "jupyter>=1.1.1",
+    "mypy>=1.16.1",
+    "pre-commit>=4.2.0",
+    "pytest>=8.4.1",
+    "pytest-cov>=6.2.1",
+    "ruff>=0.12.0",
+]
+
+[tool.uv]
+package = true
+
+[tool.uv.sources]
+torch = { index = "torch-cu126" }
+
+[[tool.uv.index]]
+name = "torch-cu126"
+url = "https://download.pytorch.org/whl/cu126"
+
+[tool.ruff]
+line-length = 120
+
+[tool.pytest.ini_options]
+addopts = "--cov src"

src/rubik/__main__.py

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

src/rubik/action.py

@@ -0,0 +1,158 @@
+import re
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from rubik.tensor_utils import build_permutation_matrix, build_cube_tensor
+
+
+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=torch.int16,
+        ),
+        # rot about Y: X -> Z
+        torch.tensor(
+            [
+                [1, 0, 0, 0],
+                [0, 0, 0, -1],
+                [0, 0, 1, 0],
+                [0, 1, 0, 0],
+            ],
+            dtype=torch.int16,
+        ),
+        # rot about Z: Y -> X
+        torch.tensor(
+            [
+                [1, 0, 0, 0],
+                [0, 0, 1, 0],
+                [0, -1, 0, 0],
+                [0, 0, 0, 1],
+            ],
+            dtype=torch.int16,
+        ),
+    ]
+)
+
+POS_SHIFTS = torch.tensor(
+    [
+        [0, 0, 0, 1],
+        [0, 1, 0, 0],
+        [0, 0, 1, 0],
+    ],
+    dtype=torch.int16,
+)
+
+
+# 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(
+    [
+        build_permutation_matrix(size=6, perm="0254"),
+        build_permutation_matrix(size=6, perm="0153"),
+        build_permutation_matrix(size=6, perm="1234"),
+    ]
+)
+
+
+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_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=torch.int16)
+
+
+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.
+    """
+    tensor = build_cube_tensor(colors=list("ULCRBD"), size=size)
+    length = 6 * (size**2)
+
+    # extract faces impacted by the move
+    indices = tensor.indices().to(dtype=torch.int16)  # 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)
+    offsets = (POS_SHIFTS[axis] * (size - 1)).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(torch.int16) @ 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)
+    inputs = inputs.transpose(0, 1).tolist()  # size = (n, 4)
+    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(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))}
+    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,
+    )
+    perm_values = torch.tensor([1] * length, dtype=torch.int16)
+    perm_size = (3, size, 2, length, length)
+    return torch.sparse_coo_tensor(indices=perm_indices, values=perm_values, size=perm_size, dtype=torch.int16)
+
+
+def parse_action_str(move: str) -> tuple[int, int, int]:
+    """
+    Convert the name of an action into a triple (axis, slice, inverse).
+    Examples:
+        'X1'  -> (0, 1, 0)
+        'X2i' -> (0, 2, 1)
+    """
+    axis = "XYZ".index(move[0])
+    slice = int(re.findall(r"^\d+", move[1:])[0])
+    inverse = int(len(move) > (1 + len(str(slice))))
+    return (axis, slice, inverse)
+
+
+def parse_actions_str(moves: str) -> list[tuple[int, int, int]]:
+    """
+    Convert a sequence of actions in a string into a list of triples (axis, slice, inverse).
+    Examples:
+        'X1 X2i'  -> [(0, 1, 0), (0, 2, 1)]
+    """
+    return [parse_action_str(move) for move in moves.strip().split()]
+
+
+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

@@ -0,0 +1,101 @@
+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.display import stringify
+from rubik.tensor_utils import build_cube_tensor
+
+
+class Cube:
+    """
+    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).
+
+    Colors filling each tensor cell are from 0 to 6, 0 being the "dark" color,
+    the rest according to order given in "colors" attribute.
+    """
+
+    def __init__(self, colors: list[str], size: int):
+        """
+        Create Cube from a given list of 6 colors and size.
+        Example:
+            cube = Cube(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
+        """
+        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
+
+    def to(self, device: str | torch.device) -> "Cube":
+        device = torch.device(device)
+        dtype = (
+            self.state.dtype
+            if self.state.device == device
+            else torch.int16
+            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}'")
+        return self
+
+    def reset_history(self) -> None:
+        """
+        Reset internal history of moves.
+        """
+        self.history = []
+        return
+
+    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
+
+    def rotate(self, moves: str) -> None:
+        """
+        Apply a sequence of moves (defined as plain string) to the cube.
+        """
+        actions = parse_actions_str(moves)
+        for action in actions:
+            self.rotate_once(*action)
+        return
+
+    def rotate_once(self, axis: int, slice: int, inverse: int) -> None:
+        """
+        Apply a move (defined as 3 coordinates) to the cube.
+        """
+        action = self.actions[axis, slice, inverse]
+        self.state = action @ self.state
+        self.history.append([axis, slice, inverse])
+        return
+
+    def compute_changes(self, moves: str) -> dict[int, int]:
+        """
+        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())
+
+    def __str__(self):
+        """
+        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)

src/rubik/display.py

@@ -0,0 +1,23 @@
+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/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.int16)
+    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.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()

tests/unit/test_action.py

@@ -0,0 +1,179 @@
+import pytest
+from typing import Iterable
+
+import torch
+
+from rubik.action import (
+    POS_ROTATIONS,
+    POS_SHIFTS,
+    FACE_ROTATIONS,
+    build_actions_tensor,
+    build_action_tensor,
+    parse_action_str,
+    parse_actions_str,
+    sample_actions_str,
+)
+
+
+def test_position_rotation_shape():
+    """
+    Test that POS_ROTATIONS has expected shape.
+    """
+    expected = (3, 4, 4)
+    observed = POS_ROTATIONS.shape
+    assert expected == observed, f"Position rotation tensor expected shape '{expected}', got '{observed}' instead"
+
+
+@pytest.mark.parametrize(
+    "axis, input, expected",
+    [
+        (0, (1, 1, 0, 0), (1, 1, 0, 0)),  # X -> X
+        (0, (1, 0, 1, 0), (1, 0, 0, -1)),  # Y -> -Z
+        (0, (1, 0, 0, 1), (1, 0, 1, 0)),  # Z -> Y
+        (1, (1, 1, 0, 0), (1, 0, 0, 1)),  # X -> Z
+        (1, (1, 0, 1, 0), (1, 0, 1, 0)),  # Y -> Y
+        (1, (1, 0, 0, 1), (1, -1, 0, 0)),  # Z -> -X
+        (2, (1, 1, 0, 0), (1, 0, -1, 0)),  # X -> -Y
+        (2, (1, 0, 1, 0), (1, 1, 0, 0)),  # Y -> X
+        (2, (1, 0, 0, 1), (1, 0, 0, 1)),  # Z -> Z
+    ],
+)
+def test_position_rotation(axis: int, input: Iterable[int], expected: Iterable[int]):
+    """
+    Test that POS_ROTATIONS behaves as expected.
+    """
+    out = POS_ROTATIONS[axis] @ torch.tensor(input, dtype=POS_ROTATIONS.dtype)
+    exp = torch.tensor(expected, dtype=POS_ROTATIONS.dtype)
+    assert torch.equal(out, exp), f"Position rotation tensor is incorrect along axis {axis}: {out} != {exp}"
+
+
+@pytest.mark.parametrize(
+    "axis, size, input, expected",
+    [
+        (0, 3, (1, 1, 1, 1), (1, 1, 1, 0)),
+        (1, 3, (1, 1, 1, 1), (1, 0, 1, 1)),
+        (2, 3, (1, 1, 1, 1), (1, 1, 0, 1)),
+    ],
+)
+def test_position_shift(axis: int, size: int, input: Iterable[int], expected: Iterable[int]):
+    """
+    Test that POS_SHIFTS behaves as expected.
+    """
+    rot = POS_ROTATIONS[axis] @ (torch.tensor(input, dtype=POS_ROTATIONS.dtype) * (size - 1))
+    out = rot + (POS_SHIFTS[axis] * (size - 1))
+    exp = torch.tensor(expected, dtype=POS_ROTATIONS.dtype) * (size - 1)
+    assert torch.equal(out, exp), f"Position shift tensor is incorrect along axis {axis}: {out} != {exp}"
+
+
+def test_face_rotation_shape():
+    """
+    Test that FACE_ROTATIONS has expected shape.
+    """
+    expected = (3, 6, 6)
+    observed = FACE_ROTATIONS.shape
+    assert expected == observed, f"Face rotation tensor expected shape '{expected}', got '{observed}' instead"
+
+
+@pytest.mark.parametrize(
+    "axis, input, expected",
+    [
+        (0, (1, 0, 0, 0, 0, 0), (0, 0, 1, 0, 0, 0)),  # rotation about X axis: 0 (Up) -> 2 (Front)
+        (1, (1, 0, 0, 0, 0, 0), (0, 1, 0, 0, 0, 0)),  # rotation about Y axis: 0 (Up) -> 1 (Left)
+        (2, (0, 1, 0, 0, 0, 0), (0, 0, 1, 0, 0, 0)),  # rotation about Z axis: 1 (Left) -> 2 (Front)
+    ],
+)
+def test_face_rotation(axis: int, input: Iterable[int], expected: Iterable[int]):
+    """
+    Test that POS_ROTATIONS behaves as expected.
+    """
+    out = torch.tensor(input, dtype=FACE_ROTATIONS.dtype) @ FACE_ROTATIONS[axis]
+    exp = torch.tensor(expected, dtype=FACE_ROTATIONS.dtype)
+    assert torch.equal(out, exp), f"Face rotation tensor is incorrect along axis {axis}: {out} != {exp}"
+
+
+@pytest.mark.parametrize("size", [2, 3, 5, 20])
+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))
+    observed = build_actions_tensor(size).shape
+    assert expected == observed, (
+        f"'build_actions_tensor' output has incorrect shape: expected shape '{expected}', got '{observed}' instead"
+    )
+
+
+@pytest.mark.parametrize(
+    "size, axis, slice, inverse",
+    [
+        (2, 2, 1, 0),
+        (3, 0, 1, 1),
+        (5, 1, 4, 0),
+    ],
+)
+def test_build_action_tensor_shape(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
+    assert expected == observed, (
+        f"'build_action_tensor' output has incorrect shape: expected shape '{expected}', got '{observed}' instead"
+    )
+
+
+@pytest.mark.parametrize(
+    "move, expected",
+    [
+        ["X1", (0, 1, 0)],
+        ["X25i", (0, 25, 1)],
+        ["Y0", (1, 0, 0)],
+        ["Y5i", (1, 5, 1)],
+        ["Z30", (2, 30, 0)],
+        ["Z512ijk", (2, 512, 1)],
+    ],
+)
+def test_parse_action_str(move: str, expected: tuple[int, int, int]):
+    """
+    Test that "parse_action_str" behaves as expected.
+    """
+    observed = parse_action_str(move)
+    assert expected == observed, (
+        f"'parse_action_str' output is incorrect: expected '{expected}', got '{observed}' instead"
+    )
+
+
+@pytest.mark.parametrize(
+    "moves, expected",
+    [
+        ["  X1 Y0 X25i Z512ijk Z30 Y5i ", [(0, 1, 0), (1, 0, 0), (0, 25, 1), (2, 512, 1), (2, 30, 0), (1, 5, 1)]],
+    ],
+)
+def test_parse_actions_str(moves: str, expected: tuple[int, int, int]):
+    """
+    Test that "parse_action_str" behaves as expected.
+    """
+    observed = parse_actions_str(moves)
+    assert expected == observed, (
+        f"'parse_actions_str' output is incorrect: expected '{expected}', got '{observed}' instead"
+    )
+
+
+@pytest.mark.parametrize(
+    "num_moves, size, seed",
+    [
+        [1, 3, 0],
+        [1, 20, 42],
+        [256, 5, 21],
+    ],
+)
+def test_sample_actions_str(num_moves: int, size: int, seed: int):
+    """
+    Test that "sample_actions_str" is deterministic and outputs parsable content.
+    """
+    moves_1 = sample_actions_str(num_moves, size, seed)
+    moves_2 = sample_actions_str(num_moves, size, seed)
+    assert moves_1 == moves_2, f"'sample_actions_str' is non-deterministic: {moves_1} != {moves_2}"
+
+    parsed = parse_actions_str(moves_1)
+    assert len(parsed) == len(moves_1.split()), "'sample_actions_str' output cannot be parsed correctly"

tests/unit/test_cube.py

@@ -0,0 +1,127 @@
+import pytest
+
+import torch
+
+from rubik.cube import Cube
+
+
+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):
+        """
+        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}"
+        )
+        assert cube.state.shape == (6 * (size**2), 7), f"'state' has incorrect shape {cube.state.shape}"
+        assert len(cube.history) == 0, "'history' field should be empty"
+
+    @pytest.mark.parametrize("device", ["cpu"])
+    def test_to(self, device: str | torch.device):
+        """
+        Test that the .to method behaves as expected.
+        """
+        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube_2 = cube.to(device)
+        assert torch.equal(cube.state, cube_2.state), "cube has different state after calling 'to' method"
+
+    def test_reset_history(self):
+        """
+        Test that the .reset_history method behaves as expected.
+        """
+        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=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"
+
+    @pytest.mark.parametrize("num_moves, seed", [[50, 42]])
+    def test_shuffle(self, num_moves: int, seed: int):
+        """
+        Test that the .shuffle method behaves as expected.
+        """
+        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube_state = cube.state.clone()
+        cube.shuffle(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"
+
+    @pytest.mark.parametrize(
+        "moves",
+        [
+            "X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i",
+            "X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i" * 2,
+        ],
+    )
+    def test_rotate(self, moves: str):
+        """
+        Test that the .rotate method behaves as expected.
+        """
+        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        cube_state = cube.state.clone()
+        cube.rotate(moves)
+        assert cube.history != [], "method 'rotate' does not update history"
+        assert not torch.equal(cube_state, cube.state), "method 'rotate' does not change state"
+
+    @pytest.mark.parametrize(
+        "axis, slice, inverse",
+        [
+            [0, 2, 0],
+            [1, 1, 1],
+            [2, 0, 0],
+        ],
+    )
+    def test_rotate_once(self, axis: int, slice: int, inverse: int):
+        """
+        Test that the .rotate_once method behaves as expected.
+        """
+        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=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 not torch.equal(cube_state, cube.state), "method 'rotate_once' does not change state"
+
+    @pytest.mark.parametrize(
+        "moves",
+        [
+            "X2 X1i Y1i",
+            "X2 X1i Y1i Z1i Y0 Z0i X2 X1i Y1i Z1i Y0 Z0i " * 2,
+        ],
+    )
+    def test_compute_changes(self, moves: str):
+        """
+        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()
+        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))]
+
+        # apply changes induced by moves using the optimized 'rotate' method
+        cube.rotate(moves)
+        observed = cube.state.argmax(dim=-1).to(torch.int16).tolist()
+
+        # assert the tow are identical
+        assert expected == observed, "method 'compute_changes' does not behave correctly: "
+
+    def test__str__(self):
+        """
+        Test that the __str__ method behaves as expected.
+        """
+        cube = Cube(colors=["U", "L", "C", "R", "B", "D"], size=3)
+        repr = str(cube)
+        assert len(repr), "__str__ method returns an empty representation"

tests/unit/test_display.py

@@ -0,0 +1,42 @@
+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

@@ -0,0 +1,38 @@
+import pytest
+
+import torch
+
+from rubik.tensor_utils import build_cube_tensor, build_permutation_matrix
+
+
+@pytest.mark.parametrize("size", [2, 3, 5, 20])
+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)
+    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()
+    z_sums = facets.sum(dim=(0, 1, 2)).tolist()
+    expected = [(size**2) + (4 * size)] + [4 * size] * (size - 2) + [(size**2) + (4 * size)]
+    assert x_sums == expected, (
+        f"'build_cube_tensor' has incorrect sum along X axis: expected '{expected}', got '{x_sums}'"
+    )
+    assert y_sums == expected, (
+        f"'build_cube_tensor' has incorrect sum along Y axis: expected '{expected}', got '{y_sums}'"
+    )
+    assert z_sums == expected, (
+        f"'build_cube_tensor' has incorrect sum along Z axis: expected '{expected}', got '{z_sums}'"
+    )
+
+
+@pytest.mark.parametrize("size, perm", [[2, "01"], [3, "210"], [6, "2345"]])
+def test_build_permutation_matrix(size: int, perm: str):
+    """
+    Test that build_permutation_matrix behaves as expected.
+    """
+    matrix = build_permutation_matrix(size, perm)
+    mapping = dict(matrix.indices().transpose(0, 1).tolist())
+    for i, j in zip(perm, perm[1:] + perm[0]):
+        assert mapping[int(i)] == int(j), f"'build_permutation_matrix' outputs has wrong behavior: {perm}, {mapping}"