def rotate(fin, rot_mat):
N_LAYERS = len(rot_mat);
FLAVORS = len(rot_mat[0]);
assert(size(fin,1) == N_LAYERS and size(fin, 2) == FLAVORS);
fout = zeros((len(fin), N_LAYERS*FLAVORS), dtype = fin.dtype);
for L in range(N_LAYERS):
fout[:, L::N_LAYERS] = fort_rot.rotate(fin[:,L,:,:], rot_mat[L]);
return fout;
def rotate_all(mat, rot_mat, need_extra = False):
Nm = len(mat);
N = len(rot_mat);
L2 = len(rot_mat[0]);
out = zeros((Nm, N*L2), dtype = mat.dtype);
for L in range(N):
mat_tmp = mat[:, L:L2*N:N, L:L2*N:N];
tmp = fort_rot.rotate(mat_tmp, rot_mat[L]);
out[:, L::N] = tmp;
if need_extra:
for n in range(N*L2, size(mat, 2)): out = c_[out, mat[:, n, n]];
return out;
