Recursion/src/recursion/dataset/common.py at simple · SimplexLab/Recursion

History

48 lines (39 loc) · 1.32 KB

Raw

from typing import List, Optional

import numpy as np

import pydantic

# Global list mapping each dihedral transform id to its inverse.

# Index corresponds to the original tid, and the value is its inverse.

DIHEDRAL_INVERSE = [0, 3, 2, 1, 4, 5, 6, 7]

class PuzzleDatasetMetadata(pydantic.BaseModel):

pad_id: int

ignore_label_id: Optional[int]

blank_identifier_id: int

vocab_size: int

seq_len: int

num_puzzle_identifiers: int

total_groups: int

mean_puzzle_examples: float

total_puzzles: int

sets: List[str]

def dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray:

"""8 dihedral symmetries by rotate, flip and mirror"""

if tid == 0:

return arr # identity

elif tid == 1:

return np.rot90(arr, k=1)

elif tid == 2:

return np.rot90(arr, k=2)

elif tid == 3:

return np.rot90(arr, k=3)

elif tid == 4:

return np.fliplr(arr) # horizontal flip

elif tid == 5:

return np.flipud(arr) # vertical flip

elif tid == 6:

return arr.T # transpose (reflection along main diagonal)

elif tid == 7:

return np.fliplr(np.rot90(arr, k=1)) # anti-diagonal reflection

else:

return arr

def inverse_dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray:

return dihedral_transform(arr, DIHEDRAL_INVERSE[tid])

Navigation Menu

Uh oh!

Uh oh!

FilesExpand file tree

common.py

Latest commit

History

common.py

File metadata and controls