def rotate(x, alpha, beta, gamma):
t = time_logging.start()
y = x.cpu().detach().numpy()
R = rot(alpha, beta, gamma)
if x.ndimension() == 4:
for i in range(y.shape[0]):
y[i] = rotate_scalar(y[i], R)
else:
y = rotate_scalar(y, R)
x = x.new_tensor(y)
if x.ndimension() == 4 and x.size(0) == 3:
rep = irr_repr(1, alpha, beta, gamma, x.dtype).to(x.device)
x = torch.einsum("ij,jxyz->ixyz", (rep, x))
time_logging.end("rotate", t)
return x
def rotate(x, alpha):
t = time_logging.start()
y = x.cpu().detach().numpy()
R = rot(alpha, 0, 0)
y = rotate_scalar(y, R)
x = x.new_tensor(y)
time_logging.end("rotate", t)
return x
def get_volumes(size=20, pad=8, rotate=False, rotate90=False):
assert size >= 4
tetris_tensorfields = [
[(0, 0, 0), (0, 0, 1), (1, 0, 0), (1, 1, 0)], # chiral_shape_1
[(0, 1, 0), (0, 1, 1), (1, 1, 0), (1, 0, 0)], # chiral_shape_2
[(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0)], # square
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3)], # line
[(0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0)], # corner
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0)], # L
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 1)], # T
[(0, 0, 0), (1, 0, 0), (1, 1, 0), (2, 1, 0)], # zigzag
]
labels = np.arange(len(tetris_tensorfields))
tetris_vox = []
for shape in tetris_tensorfields:
volume = np.zeros((4, 4, 4))
for xi_coords in shape:
volume[xi_coords] = 1
volume = zoom(volume, size / 4, order=0)
volume = np.pad(volume, pad, 'constant')
if rotate:
a, c = np.random.rand(2) * 2 * np.pi
b = np.arccos(np.random.rand() * 2 - 1)
volume = rotate_scalar(volume, rot(a, b, c))
if rotate90:
volume = rot_volume_90(volume)
tetris_vox.append(volume[np.newaxis, ...])
tetris_vox = np.stack(tetris_vox).astype(np.float32)
return tetris_vox, labels