def text_split_transform(fsplit_lst, irrep_dict, alpha, parts, mem_dict=None):
'''
fsplit_pkl: list of split file names of the distance values for a chunk of the total distance values
irrep_dict: irrep dict
alpha: weak partition
parts: list/iterable of partitions of the parts of alpha
'''
print(' Computing transform on splits: {}'.format(fsplit_lst))
cos_reps = coset_reps(sn(8), young_subgroup_perm(alpha))
save_dict = {}
cyc_irrep_func = cyclic_irreps(alpha)
pid = os.getpid()
for split_f in fsplit_lst:
with open(split_f, 'r') as f:
for line in tqdm(f):
otup, perm_tup, dist = clean_line(line)
perm_rep = irrep_dict[
perm_tup] # perm_rep is a dict of (i, j) -> matrix
block_cyclic_rep = block_cyclic_irreps(otup, cos_reps,
cyc_irrep_func)
mult_yor_block(perm_rep, dist, block_cyclic_rep, save_dict)
if mem_dict is not None:
mem_dict[pid] = max(check_memory(verbose=False),
mem_dict.get(pid, 0))
block_size = wreath_dim(parts)
n_cosets = coset_size(alpha)
mat = convert_yor_matrix(save_dict, block_size, n_cosets)
return mat
def split_transform(fsplit_lst, irrep_dict, alpha, parts, mem_dict=None):
'''
fsplit_pkl: list of pkl file names of the distance values for a chunk of the total distance values
irrep_dict: irrep dict
alpha: weak partition
parts: list/iterable of partitions of the parts of alpha
'''
print(' Computing transform on splits: {}'.format(fsplit_lst))
cos_reps = coset_reps(sn(8), young_subgroup_perm(alpha))
save_dict = {}
cyc_irrep_func = cyclic_irreps(alpha)
pid = os.getpid()
for fsplit_pkl in fsplit_lst:
with open(fsplit_pkl, 'r') as f:
# dict of function values
pkl_dict = load_pkl(fsplit_pkl)
for perm_tup, tup_dict in pkl_dict.items():
for tup, dists in tup_dict.items():
dist_tot = sum(dists)
perm_rep = irrep_dict[
perm_tup] # perm_rep is a dict of (i, j) -> matrix
block_cyclic_rep = block_cyclic_irreps(
tup, cos_reps, cyc_irrep_func)
mult_yor_block(perm_rep, dist_tot, block_cyclic_rep,
save_dict)
if mem_dict is not None:
mem_dict[pid] = max(check_memory(verbose=False),
mem_dict.get(pid, 0))
del pkl_dict
block_size = wreath_dim(parts)
n_cosets = coset_size(alpha)
mat = convert_yor_matrix(save_dict, block_size, n_cosets)
return mat
def block_pad(self, arr):
'''
Given an array of length(self.cos_reps)
Block extend this to length full group / len(self.cos_reps)
'''
block_size = wreath_dim(self.parts)**2
output = np.zeros(len(arr) * block_size)
for i in range(arr.size):
output[i * block_size:(i + 1) * block_size] = arr[i]
return output
def full_transform(args, alpha, parts, split_chunks):
print('Computing full transform for alpha: {} | parts: {}'.format(
alpha, parts))
savedir_alpha = os.path.join(args.savedir, args.alpha)
savename = os.path.join(savedir_alpha, '{}'.format(parts))
print('Savename: {}'.format(savename))
if os.path.exists(savename + '.npy'):
print('Skipping. Already computed fourier matrix for: {} | {}'.format(
alpha, parts))
exit()
manager = Manager()
irrep_dict = load_pkl(
os.path.join(args.pkldir, args.alpha, '{}.pkl'.format(parts)))
mem_usg = check_memory(verbose=False)
if not os.path.exists(savedir_alpha):
print('Making: {}'.format(savedir_alpha))
os.makedirs(savedir_alpha)
if args.par > 1:
print('Par process with {} processes...'.format(len(split_chunks)))
mem_dict = manager.dict()
with Pool(len(split_chunks)) as p:
arg_tups = [(_fn, irrep_dict, alpha, parts, mem_dict)
for _fn in split_chunks]
matrices = p.starmap(text_split_transform, arg_tups)
np.save(savename, sum(matrices))
else:
print('Single thread...')
matrices = []
block_size = wreath_dim(parts)
n_cosets = coset_size(alpha)
shape = (block_size * n_cosets, block_size * n_cosets)
result = np.zeros(shape, dtype=np.complex128)
mem_dict = {}
for _fn in split_chunks:
res = text_split_transform(_fn, irrep_dict, alpha, parts)
matrices.append(res)
result += res
np.save(savename, sum(matrices))
print('Post loading pickle mem usg: {:.4}mb | Final mem usg: {:.4f}mb'.
format(mem_usg, check_memory(False)))
print('Processes')
for pid, usg in mem_dict.items():
print('{} | {:.4f}mb'.format(pid, usg))
print('Done!')
def par_cube_ft(rank, size, alpha, parts):
start = time.time()
try:
df = load_df('/scratch/hopan/cube/')
irrep_dict = load_irrep('/scratch/hopan/cube/', alpha, parts)
except Exception as e:
print('rank {} | memory usg: {} | exception {}'.format(rank, check_memory(verbose=False), e))
print('Rank {:3d} / {} | load irrep: {:.2f}s | mem: {}mb'.format(rank, size, time.time() - start, check_memory(verbose=False)))
cos_reps = coset_reps(sn(8), young_subgroup_perm(alpha))
save_dict = {}
cyc_irrep_func = cyclic_irreps(alpha)
chunk_size = len(df) // size
start_idx = chunk_size * rank
#print('Rank {} | {:7d}-{:7d}'.format(rank, start_idx, start_idx + chunk_size))
if rank == 0:
print('Rank {} | elapsed: {:.2f}s | {:.2f}mb | done load'.format(rank, time.time() - start, check_memory(verbose=False)))
for idx in range(start_idx, start_idx + chunk_size):
row = df.loc[idx]
otup = tuple(int(i) for i in row[0])
perm_tup = tuple(int(i) for i in row[1])
dist = int(row[2])
perm_rep = irrep_dict[perm_tup] # perm_rep is a dict of (i, j) -> matrix
block_cyclic_rep = block_cyclic_irreps(otup, cos_reps, cyc_irrep_func)
mult_yor_block(perm_rep, dist, block_cyclic_rep, save_dict)
if rank == 0:
print('Rank {} | elapsed: {:.2f}s | {:.2f}mb | done add'.format(rank, time.time() - start, check_memory(verbose=False)))
del irrep_dict
block_size = wreath_dim(parts)
n_cosets = coset_size(alpha)
mat = convert_yor_matrix(save_dict, block_size, n_cosets)
if rank == 0:
print('Rank {} | elapsed: {:.2f}s | {:.2f}mb | done matrix conversion'.format(rank, time.time() - start, check_memory(verbose=False)))
return mat
