ob_master = {
'J1': J1,
'J2': J2,
'phi1': phi1,
'phi2': phi2,
'tau': zeta,
'turn': turn
}
return ob_master
suffixes = ['_init.h5', '_first.h5', '_skip.h5', '_last.h5']
ob_funcs = [get_ob_init, get_ob_first, get_ob_skip, get_ob_last]
for filelist, fname in zip(filelists, fnames):
for suffix, get_ob in zip(suffixes, ob_funcs):
fname_suff = fname + suffix
print(fname_suff)
outfile = out_folder + fname_suff
obs = []
for myfile in filelist:
obs.append(get_ob(in_folder + myfile))
master_ob = merge_obs(obs)
optics = kfm.h5_to_dict(myfile, group='beam-optics')
particle_on_CO = kfm.h5_to_dict(myfile, group='closed-orbit')
kfm.dict_to_h5(master_ob, outfile, group='data', readwrite_opts='w')
kfm.dict_to_h5(optics, outfile, group='optics', readwrite_opts='a')
kfm.dict_to_h5(particle_on_CO,
outfile,
group='closed-orbit',
readwrite_opts='a')
J1, J2 = distribution.J1_J2_from_physical(init_denormalized_6D, optics['invW'], partCO)
ecloud_lattice = ec_stl.LatticeWithEclouds(line, eclouds_info, ps, device=device)
ecloud_lattice.set_optics_CO(optics, partCO)
ecloud_lattice.add_tricub_data(sys.argv[2], 'drift', max_z=0.15)
#ecloud_lattice.add_tricub_data('refined_Pinch10_MTI4.0_MLI2.0_DTO1.0_DLO1.0.h5', 'drift', max_z=0.15)
tricub_to_tricub_data = {}
for key in eclouds_info['length'].keys():
tricub_to_tricub_data[key] = 'drift'
ecloud_lattice.finalize_assignments(tricub_to_tricub_data)
end_setup_time = time.time()
print(f'Setting up time: {(end_setup_time - start_time)/60.}mins')
ecloud_lattice.fma_tracking(distance_between_tunes = 1000, until_turn = 20000, num_stores=n_stores)
output_to_save = {'turn_q' : np.array(ecloud_lattice.turn_q_list),
'tune_is_valid' : np.array(ecloud_lattice.tune_is_valid_list),
'q1' : np.array(ecloud_lattice.q1_list),
'q2' : np.array(ecloud_lattice.q2_list),
'qx' : np.array(ecloud_lattice.qx_list),
'qy' : np.array(ecloud_lattice.qy_list),
'J1' : J1/se1**2,
'J2' : J2/se2**2
}
#kfm.dict_to_h5(output_to_save, f'fma_tunes_ec_scale{ecloud_scale:.2f}_IMO_0_ref.h5')
kfm.dict_to_h5(output_to_save, f'fma_tunes_ec_scale{ecloud_scale:.2f}_IMO_0_'+sys.argv[2]+'.h5')
N_pinches=n_pinches_to_average,
save_sigmas_and_coords=True,
grid=grid_dict,
eden=eden)
with multiprocessing.Pool(workers) as pool:
result_list = pool.starmap_async(
kern, [(ii, result_dict, lock, eden)
for ii in range(n_pinches_to_average - 1)])
print(result_list.get())
# Save pinch to file
out_pinch = pinch_name + '.h5'
kfm.dict_to_h5(grid_dict, out_pinch, group='grid', readwrite_opts='w')
for i in range(result_dict['phi'].shape[0]):
kfm.dict_to_h5(
{
'phi': result_dict['phi'][i, :, :],
'rho': result_dict['rho'][i, :, :]
},
out_pinch,
group='slices/slice%d' % i,
readwrite_opts='a')
kfm.dict_to_h5({'ti_method': 'FD'},
out_pinch,
group='stats',
readwrite_opts='a')
kfm.dict_to_h5({'sim': str(pp_dict)},
'J1': J1.reshape(shape) / se1**2,
'J2': J2.reshape(shape) / se2**2
}
last_dict = {
'x': parts.x.reshape(shape),
'px': parts.px.reshape(shape),
'y': parts.y.reshape(shape),
'py': parts.py.reshape(shape),
'zeta': parts.zeta.reshape(shape),
'delta': parts.delta.reshape(shape),
'at_turn': parts.at_turn.reshape(shape)
}
time_dict = {
'setup_time_mins': (end_setup_time - start_time) / 60.,
'tracking_time_mins': (end_tracking - start_tracking) / 60.
}
kfm.dict_to_h5(init_dict, output_file, group='input', readwrite_opts='w')
kfm.dict_to_h5(last_dict, output_file, group='output', readwrite_opts='a')
kfm.dict_to_h5(partCO, output_file, group='closed-orbit', readwrite_opts='a')
kfm.dict_to_h5(optics, output_file, group='optics', readwrite_opts='a')
#args_dict = vars(args)
#for key in args_dict.keys():
# if args_dict[key] is None:
# args_dict[key] = 'None'
#kfm.dict_to_h5(args_dict, output_file, group='args', readwrite_opts='a')
kfm.dict_to_h5(time_dict, output_file, group='time', readwrite_opts='a')
Nx = ob['phi'].shape[0]
Ny = ob['phi'].shape[1]
Nz = ob['phi'].shape[2]
N_keep_x = 101
N_keep_y = 101
N_keep_z = 200
x0 = (Nx - N_keep_x) // 2
y0 = (Ny - N_keep_y) // 2
z0 = (Nz - N_keep_z) // 2
dict_new_file = {}
dict_new_file['phi'] = ob['phi'][x0:-x0, y0:-y0, z0:-z0]
dict_new_file['rho'] = ob['rho'][x0:-x0, y0:-y0, z0:-z0]
dict_new_file['xg'] = ob['xg'][x0:-x0]
dict_new_file['yg'] = ob['yg'][y0:-y0]
dict_new_file['zg'] = ob['zg'][z0:-z0]
print('Old shape: ', ob['phi'].shape)
print('New shape: ', dict_new_file['phi'].shape)
print('x, y, z new Ranges:')
print(dict_new_file['xg'][0], dict_new_file['xg'][-1])
print(dict_new_file['yg'][0], dict_new_file['yg'][-1])
print(dict_new_file['zg'][0], dict_new_file['zg'][-1])
kfm.dict_to_h5(dict_new_file,
pinch_folder + pinch + '_cut.h5',
compression_opts=0)
'Nx': len(xg_new),
'Ny': len(yg_new),
'Nz': len(zg_new)
}
settings_dict = {
'magnify_transverse_in': magnify_transverse_in,
'magnify_longitudinal_in': magnify_longitudinal_in,
'demagnify_transverse_out': demagnify_transverse_out,
'demagnify_longitudinal_out': demagnify_longitudinal_out,
'symmetric2D': do_symmetric2D
}
kfm.dict_to_h5(grid_dict,
out_fname,
compression_opts=compression_opts,
group='grid',
readwrite_opts='w')
kfm.dict_to_h5(settings_dict,
out_fname,
compression_opts=compression_opts,
group='settings',
readwrite_opts='a')
kfm.dict_to_h5(grid_dict,
out_efname,
compression_opts=compression_opts,
group='grid',
readwrite_opts='w')
kfm.dict_to_h5(settings_dict,
out_efname,
compression_opts=compression_opts,
import sys
sys.path.append('../Tools')
import kostas_filemanager as kfm
list_of_pinches = ['/eos/user/k/kparasch/Pinches/ArcDipole_07.h5',
'/eos/user/k/kparasch/Pinches/ArcDipole_07_2000.h5'
]
numbers_of_pinches = [4000,2000]
temp_file = 'temp.h5'
out_file = '/eos/user/e/ecincohe/Raw_Pinches/LHC_ArcDip_1.35sey_0.7e11ppb.h5'
fname0 = list_of_pinches[0]
grid_dict = kfm.h5_to_dict(fname0, group='grid')
stats_dict = kfm.h5_to_dict(fname0, group='stats')
kfm.dict_to_h5(grid_dict, temp_file, group='grid', readwrite_opts='w')
kfm.dict_to_h5(stats_dict, temp_file, group='stats', readwrite_opts='a')
nx = len(grid_dict['xg'])
ny = len(grid_dict['yg'])
result_dict = {'phi': np.zeros([nx,ny]), 'rho': np.zeros([nx,ny])}
n_slices = len(grid_dict['zg'])
total_pinches = sum(numbers_of_pinches)
for kk in range(n_slices):
print(f'Slice: {kk}/{n_slices}')
this_slice = f'slices/slice{kk}'
phi = np.zeros([nx,ny])