def add_tricub_data(self, fname, data_name, max_z=None):
grid = kfm.h5_to_obj(fname, group='grid')
if max_z is not None:
slices = np.where(abs(grid.zg) < max_z)[0]
nz = len(slices)
z0 = grid.zg[slices[0]]
else:
nz = grid.Nz
slices = range(nz)
z0 = grid.z0
nx = len(grid.xg)
ny = len(grid.yg)
self.tricub_data_indices[data_name] = self.tricub_data_buffer.n_objects
self.tricub_data[data_name] = sixtracklib.TriCubData(
cbuffer=self.tricub_data_buffer, nx=nx, ny=ny, nz=nz)
tricub_data = self.tricub_data[data_name]
tricub_data.x0 = grid.x0
tricub_data.y0 = grid.y0
tricub_data.z0 = z0
tricub_data.dx = grid.dx
tricub_data.dy = grid.dy
tricub_data.dz = grid.dz
if kfm.h5_to_dict(fname, group='settings')['symmetric2D']:
tricub_data.mirror_x = 1
tricub_data.mirror_y = 1
tricub_data.mirror_z = 0
else:
tricub_data.mirror_x = 0
tricub_data.mirror_y = 0
tricub_data.mirror_z = 0
scale = [
1., grid.dx, grid.dy, grid.dz, grid.dx * grid.dy,
grid.dx * grid.dz, grid.dy * grid.dz, (grid.dx * grid.dy) * grid.dz
]
for kk in slices:
kk0 = kk - slices[0]
if kk0 % 10 == 0:
print(f'Loading {data_name} ecloud data ({kk0}/{nz})')
phi = kfm.h5_to_dict(fname, group=f'slices/slice{kk}')['phi']
for ll in range(8):
phi[:, :, ll] *= scale[ll]
index = 8 * (nx * (ny * kk0))
tricub_data.table_addr[index:index +
8 * nx * ny] = phi[:, :, :].transpose(
1, 0, 2).flatten()
self.tricub_data_buffer_ids[data_name] = self.job.add_stored_buffer(
cbuffer=self.tricub_data_buffer)
return
def get_ob_first(myfile):
print(myfile)
out_data = kfm.h5_to_dict(myfile, group='output')
optics = kfm.h5_to_dict(myfile, group='beam-optics')
particle_on_CO = kfm.h5_to_dict(myfile, group='closed-orbit')
e1 = optics['epsn_x'] / (optics['beta0'] * optics['gamma0'])
e2 = optics['epsn_y'] / (optics['beta0'] * optics['gamma0'])
invW = optics['invW']
part_on_CO = np.array([
particle_on_CO['x'], particle_on_CO['px'], particle_on_CO['y'],
particle_on_CO['py'], particle_on_CO['zeta'], particle_on_CO['delta']
]).T
X_first = np.array([
out_data['x_tbt_first'].T, out_data['px_tbt_first'].T,
out_data['y_tbt_first'].T, out_data['py_tbt_first'].T,
out_data['zeta_tbt_first'].T, out_data['delta_tbt_first'].T
]).T
X_first_norm = np.tensordot(X_first - part_on_CO, invW, axes=(-1, 1))
J1_first = 0.5 * (X_first_norm[:, :, 0]**2 + X_first_norm[:, :, 1]**2) / e1
J2_first = 0.5 * (X_first_norm[:, :, 2]**2 + X_first_norm[:, :, 3]**2) / e2
phi1_first = np.arctan2(-X_first_norm[:, :, 1], X_first_norm[:, :, 0])
phi2_first = np.arctan2(-X_first_norm[:, :, 3], X_first_norm[:, :, 2])
zeta_first = X_first[:, :, 4]
turn_first = out_data['at_turn_tbt_first']
ob = {
'J1': J1_first,
'J2': J2_first,
'phi1': phi1_first,
'phi2': phi2_first,
'tau': zeta_first,
'turn': turn_first
}
return ob
def get_ob_input(myfile):
print(myfile)
#out_data = kfm.h5_to_dict(myfile, group = 'output')
in_data = kfm.h5_to_dict(myfile, group='input')
optics = kfm.h5_to_dict(myfile, group='beam-optics')
particle_on_CO = kfm.h5_to_dict(myfile, group='closed-orbit')
e1 = optics['epsn_x'] / (optics['beta0'] * optics['gamma0'])
e2 = optics['epsn_y'] / (optics['beta0'] * optics['gamma0'])
invW = optics['invW']
part_on_CO = np.array([
particle_on_CO['x'], particle_on_CO['px'], particle_on_CO['y'],
particle_on_CO['py'], particle_on_CO['zeta'], particle_on_CO['delta']
]).T
X_init = np.array([
in_data['init_x'].T, in_data['init_px'].T, in_data['init_y'].T,
in_data['init_py'].T, in_data['init_zeta'].T, in_data['init_delta'].T
]).T
X_init = np.array([X_init])
X_init_norm = np.tensordot(X_init, invW, axes=(-1, 1))
J1_init = 0.5 * (X_init_norm[:, :, 0]**2 + X_init_norm[:, :, 1]**2) / e1
J2_init = 0.5 * (X_init_norm[:, :, 2]**2 + X_init_norm[:, :, 3]**2) / e2
phi1_init = np.arctan2(-X_init_norm[:, :, 1], X_init_norm[:, :, 0])
phi2_init = np.arctan2(-X_init_norm[:, :, 3], X_init_norm[:, :, 2])
zeta_init = X_init[:, :, 4]
turn_init = np.zeros_like(J1_init)
ob = {
'J1': J1_init,
'J2': J2_init,
'phi1': phi1_init,
'phi2': phi2_init,
'tau': zeta_init,
'turn': turn_init
}
return ob
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')
from TricubicInterpolation import cTricubic as ti
import numpy as np
import matplotlib.pyplot as plt
import sys
#sys.path.append('..')
import kostas_filemanager as kfm
plt.style.use('kostas')
pinch = 'eclouds/refined_Pinch7_cut_MTI4.0_MLI2.0_DTO1.0_DLO1.0.h5'
epinch = 'eclouds/e_refined_Pinch7_cut_MTI4.0_MLI2.0_DTO1.0_DLO1.0.h5'
grid = kfm.h5_to_dict(pinch, group='grid')
Nx, Ny, Nz = grid['Nx'], grid['Ny'], grid['Nz']
dx, dy, dz = grid['dx'], grid['dy'], grid['dz']
x0, y0, z0 = grid['x0'], grid['y0'], grid['z0']
xg, yg, zg = grid['xg'], grid['yg'], grid['zg']
phi = np.empty([Nx, Ny, Nz, 8])
ex = np.empty([Nx, Ny, Nz, 8])
ey = np.empty([Nx, Ny, Nz, 8])
ez = np.empty([Nx, Ny, Nz, 8])
for kk in range(Nz):
phi[:, :, kk, :] = kfm.h5_to_dict(pinch,
group='slices/slice%d' % kk)['phi']
ex[:, :, kk, :] = kfm.h5_to_dict(epinch,
group='slices/ex_slice%d' % kk)['ex']
import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio
import sys
#sys.path.append('..')
import kostas_filemanager as kfm
pinch_folder = '/home/kparasch/Storage/MyPinches/'
pinch = 'Pinch7'
print('Loading file...')
ob = kfm.h5_to_dict(pinch_folder + pinch + '.h5')
print('Shape of xg', ob['xg'].shape)
print('Shape of yg', ob['yg'].shape)
print('Shape of zg', ob['zg'].shape)
print('Shape of phi', ob['phi'].shape)
i_zero = np.argmin(np.abs(ob['xg']))
j_zero = np.argmin(np.abs(ob['yg']))
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
perc = 0.
print('Reading... ')
ti_method = 'Exact-Mirror2' if do_symmetric2D else 'Exact'
plot_hist_str = 'plt.hist(bins[:-1], bins, weights=counts, histtype=\'step\', log=True)'
hist_counts, bins = np.histogram([], bins=2000, range=(-13, 0))
log10_vx_hist = np.zeros_like(hist_counts)
log10_vy_hist = np.zeros_like(hist_counts)
log10_vz_hist = np.zeros_like(hist_counts)
max_log10_vx = -20
max_log10_vy = -20
max_log10_vz = -20
kk = 0
grid = kfm.h5_to_dict(out_fname, group='grid')
phi0 = kfm.h5_to_dict(out_fname, group='slices/slice%d' % kk)['phi']
ex0 = kfm.h5_to_dict(out_efname, group='slices/ex_slice%d' % kk)['ex']
ey0 = kfm.h5_to_dict(out_efname, group='slices/ey_slice%d' % kk)['ey']
ez0 = kfm.h5_to_dict(out_efname, group='slices/ez_slice%d' % kk)['ez']
for kk in range(1, len(zg_new)):
if (1. * kk) / len(zg_new) >= perc:
print('%d%%... ' % (int(kk / slice_index * 100)))
perc += 0.1
phi1 = kfm.h5_to_dict(out_fname, group='slices/slice%d' % kk)['phi']
ex1 = kfm.h5_to_dict(out_efname, group='slices/ex_slice%d' % kk)['ex']
ey1 = kfm.h5_to_dict(out_efname, group='slices/ey_slice%d' % kk)['ey']
ez1 = kfm.h5_to_dict(out_efname, group='slices/ez_slice%d' % kk)['ez']
TIphi = pyTI.Tricubic_Interpolation(A=np.array([phi0, phi1
]).transpose(1, 2, 0, 3),
dx=grid['dx'],
import shutil
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}')