Spaces:
Sleeping
Sleeping
| from dataclasses import dataclass | |
| import torch | |
| class Cube: | |
| """ | |
| A 5D tensor filled with 0 or 1. Dimensions have the following interpretation: | |
| - 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). | |
| """ | |
| state: torch.Tensor | |
| colors: list[str] | |
| size: int | |
| def from_default(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) | |
| """ | |
| 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 | |
| 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) | |
| 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] | |
| def to_grid(self, pad_colors: bool = False) -> list[list[list[str]]]: | |
| """ | |
| Convert Cube into a 3D grid representation. | |
| """ | |
| 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] | |
| 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) | |
| 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]) | |
| return "\n".join([l1, l2, l3]) | |