geometry.py
13.6 kB · python · 388 lines
1import numpy as np2from enum import Enum3from typing import List, Optional, Tuple4from mrlypy.core.errors import MrlyError5from mrlypy.two.models import Cell2d6from mrlypy.three.models import Cell3d78class Orientation(Enum):9 HORIZONTAL = "horizontal"10 VERTICAL = "vertical"1112# CONSTANTS1314VOID = 015FILL = 116GRID = 217UP = 318LEFT = 419RIGHT = 52021# HELPERS2223def is_cube(cell: Cell3d) -> bool:24 return cell.width == cell.height == cell.depth2526def is_hex(cell) -> bool:27 h, w = cell.height, cell.width28 if w > h:29 if w % 2 == 0:30 return False31 dx = (3 * (w + 1)) // 432 row_shift = h // 233 if (dx + row_shift) % 2 != 0:34 return False35 return True36 elif h > w:37 dy = (3 * (h + 1)) // 438 row_shift = w // 239 if (dy + row_shift) % 2 != 0:40 return False41 return True42 return False4344def get_orientation(width, height) -> Orientation:45 if width > height:46 return Orientation.HORIZONTAL47 if height > width:48 return Orientation.VERTICAL49 raise MrlyError("Cell must be a hexagon.")5051def check_orientation(orientation) -> Orientation:52 if isinstance(orientation, Orientation):53 return orientation54 if orientation not in [Orientation.HORIZONTAL.value, Orientation.VERTICAL.value]:55 raise MrlyError("Unknown orientation.")56 return Orientation(orientation)5758def blank(radius: int, orientation: str, fill: int = 1, void: int = 0) -> Cell2d:59 orientation = check_orientation(orientation)60 n = radius61 match orientation:62 case Orientation.HORIZONTAL:63 height = 2 * n64 width = 4 * n - 165 types = np.full((height, width), fill, dtype=np.uint8)66 for r in range(height):67 p = max(0, n - 1 - r, r - n)68 if p > 0:69 types[r, :p] = void70 types[r, width-p:] = void71 case Orientation.VERTICAL:72 width = 2 * n73 height = (7 * n - 1) // 274 row_shift = width // 275 while ((3 * (height + 1)) // 4 + row_shift) % 2 != 0:76 height += 177 types = np.full((height, width), fill, dtype=np.uint8)78 for r in range(height):79 p = max(0, n - 1 - r, r - (height - n))80 if p > 0:81 types[r, :p] = void82 types[r, width-p:] = void83 return Cell2d(types=types)8485def pad(cell, k: int = 1, val: int = 0):86 if k < 1:87 return cell88 inner = cell._cell if hasattr(cell, '_cell') else cell89 if not is_hex(inner):90 raise MrlyError("Cell must be a hexagon.")91 orientation = get_orientation(inner.width, inner.height)92 match orientation:93 case Orientation.HORIZONTAL:94 n = inner.height // 295 case Orientation.VERTICAL:96 n = inner.width // 297 n_new = n + k98 base = blank(n_new, orientation, fill=val, void=GRID)99 y_off = (base.height - inner.height) // 2100 x_off = (base.width - inner.width) // 2101 src_types = inner.types.copy()102 src_types[src_types == GRID] = val103 h_paste = min(inner.height, base.height - y_off)104 w_paste = min(inner.width, base.width - x_off)105 base.types[y_off:y_off+h_paste, x_off:x_off+w_paste] = src_types[:h_paste, :w_paste]106 if inner._colors is not None:107 base.colors = np.zeros((base.height, base.width, 4), dtype=np.uint8)108 base.colors[y_off:y_off+h_paste, x_off:x_off+w_paste] = inner.colors[:h_paste, :w_paste]109 if inner._tags is not None:110 base.tags = np.full((base.height, base.width), val, dtype=np.uint8)111 base.tags[y_off:y_off+h_paste, x_off:x_off+w_paste] = inner.tags[:h_paste, :w_paste]112 return base113114# GEOMETRY115116def iso(cell: Cell3d):117 from .models import Cell6d118 if not is_cube(cell):119 raise MrlyError("Cell must be a cube.")120 grid = cell.types121 n_x, n_y, n_z = grid.shape122 N = n_x123 width = 2 * N124 height = 4 * N - 1125 types = np.full((height, width), GRID, dtype=np.uint8)126 for z in range(n_z):127 for y in range(n_y):128 for x in range(n_x):129 if grid[x, y, z]:130 gx = x - y + (N - 1)131 gy = x + y - 2 * z + (2 * N - 2)132 if 0 <= gx < width - 1 and 0 <= gy < height - 2:133 types[gy, gx] = UP134 types[gy, gx + 1] = UP135 types[gy + 1, gx] = LEFT136 types[gy + 1, gx + 1] = RIGHT137 types[gy + 2, gx] = LEFT138 types[gy + 2, gx + 1] = RIGHT139 return Cell6d(cell=Cell2d(types=types), projection="iso", orientation="vertical", start=1)140141def pro(cell: Cell3d):142 from .models import Cell6d143 if not is_cube(cell):144 raise MrlyError("Cell must be a cube.")145 grid = cell.types146 n_x, n_y, n_z = grid.shape147 N = n_x148 width = 2 * N149 height = 4 * N - 1150 types = np.full((height, width), GRID, dtype=np.uint8)151 y = n_y - 1152 for z in range(n_z):153 for x in range(n_x):154 val = grid[x, y, z]155 gx = x - y + (N - 1)156 gy = x + y - 2 * z + (2 * N - 2)157 draw_val = FILL if val == 1 else VOID158 if 0 <= gx < width - 1 and 0 <= gy < height - 2:159 types[gy+1, gx] = draw_val160 types[gy+2, gx] = draw_val161 x = n_x - 1162 for z in range(n_z):163 for y in range(n_y):164 val = grid[x, y, z]165 gx = x - y + (N - 1)166 gy = x + y - 2 * z + (2 * N - 2)167 draw_val = FILL if val == 1 else VOID168 if 0 <= gx < width - 1 and 0 <= gy < height - 2:169 types[gy+1, gx+1] = draw_val170 types[gy+2, gx+1] = draw_val171 z = n_z - 1172 for y in range(n_y):173 for x in range(n_x):174 val = grid[x, y, z]175 gx = x - y + (N - 1)176 gy = x + y - 2 * z + (2 * N - 2)177 draw_val = FILL if val == 1 else VOID178 if 0 <= gx < width - 1 and 0 <= gy < height - 2:179 types[gy, gx] = draw_val180 types[gy, gx+1] = draw_val181 return Cell6d(cell=Cell2d(types=types), projection="pro", orientation="vertical", start=1)182183def cut(cell: Cell3d):184 from .models import Cell6d185 if not is_cube(cell):186 raise MrlyError("Cell must be a cube.")187 scale = 4188 grid = cell.types189 block = np.ones((scale, scale, scale), dtype=np.uint8)190 grid = np.kron(grid, block).astype(np.uint8)191 size = grid.shape[0]192 k = (3 * (size - 1)) // 2193 rows = []194 for z in range(0, size, 2):195 target = k - z196 min_x = max(0, target - (size - 1))197 max_x = min(size - 1, target)198 if min_x > max_x:199 continue200 row_bits = []201 for x in range(min_x, max_x + 1):202 y = target - x203 val = grid[x, y, z]204 row_bits.append(str(val))205 rows.append("".join(row_bits))206 if not rows:207 return Cell6d(cell=Cell2d(width=1, height=1), projection="cut", orientation="horizontal", start=0)208 width = max(len(row) for row in rows)209 height = len(rows)210 types = np.full((height, width), GRID, dtype=np.uint8)211 for r, row in enumerate(rows):212 padding_total = width - len(row)213 offset = padding_total // 2214 for c, char in enumerate(row):215 if char == '1':216 types[r, c + offset] = FILL217 elif char == '0':218 types[r, c + offset] = VOID219 return Cell6d(cell=Cell2d(types=types), projection="cut", orientation="horizontal", start=0)220221# TILING222223def tessellate(cell, mask: np.ndarray):224 inner = cell._cell if hasattr(cell, '_cell') else cell225 if not is_hex(inner):226 raise MrlyError("Cell must be a hexagon.")227 orientation = get_orientation(inner.width, inner.height)228 tile_h, tile_w = inner.height, inner.width229 match orientation:230 case Orientation.HORIZONTAL:231 dx = (3 * (tile_w + 1)) // 4232 dy = tile_h233 row_shift = tile_h // 2234 case Orientation.VERTICAL:235 dx = tile_w236 dy = (3 * (tile_h + 1)) // 4237 row_shift = tile_w // 2238 positions = []239 rows, cols = np.nonzero(mask)240 if len(rows) == 0:241 return Cell2d(width=1, height=1)242 for r, c in zip(rows, cols):243 match orientation:244 case Orientation.HORIZONTAL:245 pos_x = c * dx246 pos_y = r * dy247 if c % 2 != 0:248 pos_y += row_shift249 case Orientation.VERTICAL:250 pos_x = c * dx251 pos_y = r * dy252 if r % 2 != 0:253 pos_x += row_shift254 positions.append((pos_x, pos_y))255 min_x = min(p[0] for p in positions)256 min_y = min(p[1] for p in positions)257 max_x = max(p[0] + tile_w for p in positions)258 max_y = max(p[1] + tile_h for p in positions)259 final_w = max_x - min_x260 final_h = max_y - min_y261 bg_val = GRID262 new_types = np.full((final_h, final_w), bg_val, dtype=np.uint8)263 new_colors = None264 if inner._colors is not None:265 new_colors = np.zeros((final_h, final_w, 4), dtype=np.uint8)266 new_tags = None267 if inner._tags is not None:268 new_tags = np.zeros((final_h, final_w), dtype=np.uint8)269 for (r, c), (px, py) in zip(zip(rows, cols), positions):270 dest_x = px - min_x271 dest_y = py - min_y272 src_types = inner.types273 target_slice_types = new_types[dest_y:dest_y+tile_h, dest_x:dest_x+tile_w]274 mask_paste = (src_types != bg_val)275 target_slice_types[mask_paste] = src_types[mask_paste]276 if new_colors is not None:277 src_colors = inner.colors278 target_slice_colors = new_colors[dest_y:dest_y+tile_h, dest_x:dest_x+tile_w]279 target_slice_colors[mask_paste] = src_colors[mask_paste]280 if new_tags is not None:281 src_tags = inner.tags282 target_slice_tags = new_tags[dest_y:dest_y+tile_h, dest_x:dest_x+tile_w]283 target_slice_tags[mask_paste] = src_tags[mask_paste]284 return Cell2d(types=new_types, colors=new_colors, tags=new_tags)285286# TILE287288def get_tile_mask(width: int, height: int) -> np.ndarray:289 return np.ones((height, width), dtype=np.uint8)290291def tile(cell, width: int, height: int):292 mask = get_tile_mask(width, height)293 return tessellate(cell, mask)294295def tile_crop(cell, size: Tuple[int, int]):296 inner = cell._cell if hasattr(cell, '_cell') else cell297 w, h = size298 orientation = get_orientation(w, h)299 match orientation:300 case Orientation.HORIZONTAL:301 crop_x = (w - 1) // 4302 crop_y = h // 2303 case Orientation.VERTICAL:304 crop_x = w // 2305 crop_y = (h - 1) // 4306 current_h, current_w = inner.types.shape307 start_y = crop_y308 end_y = current_h - crop_y309 start_x = crop_x310 end_x = current_w - crop_x311 if start_y >= end_y or start_x >= end_x:312 return Cell2d(types=np.zeros((1, 1), dtype=np.uint8))313 new_types = inner.types[start_y:end_y, start_x:end_x]314 new_colors = None315 if inner._colors is not None:316 new_colors = inner.colors[start_y:end_y, start_x:end_x]317 new_tags = None318 if inner._tags is not None:319 new_tags = inner.tags[start_y:end_y, start_x:end_x]320 return Cell2d(types=new_types, colors=new_colors, tags=new_tags)321322# RADIAL323324def get_radial_mask(radius: int, orientation: str) -> np.ndarray:325 if radius < 1:326 return np.zeros((1, 1), dtype=np.uint8)327 orientation = check_orientation(orientation)328 size = 2 * radius - 1329 center = radius - 1330 mask = np.zeros((size, size), dtype=np.uint8)331 match orientation:332 case Orientation.HORIZONTAL:333 c_q = center334 c_r = center - (center - (center & 1)) // 2335 case Orientation.VERTICAL:336 c_q = center - (center - (center & 1)) // 2337 c_r = center338 for r in range(size):339 for c in range(size):340 match orientation:341 case Orientation.HORIZONTAL:342 q = c343 r_axial = r - (c - (c & 1)) // 2344 case Orientation.VERTICAL:345 q = c - (r - (r & 1)) // 2346 r_axial = r347 dq = q - c_q348 dr = r_axial - c_r349 if (abs(dq) + abs(dr) + abs(dq + dr)) / 2 < radius:350 mask[r, c] = 1351 return mask352353def radial(cell, radius: int):354 inner = cell._cell if hasattr(cell, '_cell') else cell355 if not is_hex(inner):356 raise MrlyError("Cell must be a hexagon.")357 orientation = get_orientation(inner.width, inner.height)358 mask = get_radial_mask(radius, orientation)359 return tessellate(cell, mask)360361def radial_crop(cell, radius: int, size: Tuple[int, int]):362 inner = cell._cell if hasattr(cell, '_cell') else cell363 w, h = size364 orientation = get_orientation(w, h)365 match orientation:366 case Orientation.HORIZONTAL:367 row_shift = h // 2368 crop_x = row_shift369 crop_y = (radius - 1) * row_shift370 case Orientation.VERTICAL:371 row_shift = w // 2372 crop_y = row_shift373 crop_x = (radius - 1) * row_shift374 current_h, current_w = inner.types.shape375 start_y = crop_y376 end_y = current_h - crop_y377 start_x = crop_x378 end_x = current_w - crop_x379 if start_y >= end_y or start_x >= end_x:380 return Cell2d(types=np.zeros((1, 1), dtype=np.uint8))381 new_types = inner.types[start_y:end_y, start_x:end_x]382 new_colors = None383 if inner._colors is not None:384 new_colors = inner.colors[start_y:end_y, start_x:end_x]385 new_tags = None386 if inner._tags is not None:387 new_tags = inner.tags[start_y:end_y, start_x:end_x]388 return Cell2d(types=new_types, colors=new_colors, tags=new_tags)