Newer
Older
Output: x (Nelem,Nptcl,7)
For binary characters see https://docs.python.org/2/library/struct.html
2014.10.06
file.read(1)
file.read(1) # 125,100
Nelem = unpack("i", file.read(4))[0]
Nptcl = unpack("i", file.read(4))[0]
Ibeam = unpack("d", file.read(8))[0]
freq = unpack("d", file.read(8))[0]
mass = unpack("d", file.read(8))[0]
x = []
i_unk = []
i_elem = []
s = []
phs = []
ken = []
i_unk.append(unpack("b", file.read(1))[0])
i_elem.append(unpack("i", file.read(4))[0])
s.append(unpack("d", file.read(8))[0])
phs.append(unpack("d", file.read(8))[0])
ken.append(unpack("d", file.read(8))[0])
x.append(
[[unpack("f", file.read(4))[0] for l in range(7)] for k in range(Nptcl)]
)
return x, mass, freq, Ibeam, i_unk, i_elem, s, phs, ken
def x2plt(x, mass, freq, Ibeam, i_unk, i_elem, s, phs, ken, path_file="dtl1_new.plt"):
"""
Output a TraceWin's .plt file from x and etc.
Input: x (Nelem,Nptcl,7)
For binary characters see https://docs.python.org/2/library/struct.html
"""
fname = open(path_file, "w")
out = pack("b", 125)
out += pack("b", 100)
out += pack("i", len(x))
out += pack("i", len(x[0]))
out += pack("d", Ibeam)
out += pack("d", freq)
out += pack("d", mass)
out += pack("b", i_unk[i])
# out+=pack('i',i_elem[i])
out += pack("i", i) # To truncate the elem number
out += pack("d", s[i])
out += pack("d", phs[i])
out += pack("d", ken[i])
x_i = list(chain(*x[i]))
for x_ik in x_i:
out += pack("f", x_ik)
fname.write(out + "\n")
Calculate twiss from x. Note sig = sqrt(bet*eps).
Input : x (Nptcl,6)
Output: eps,bet,alf,gam
2015.07.30
x = [x_i - numpy.mean(x_i) for x_i in numpy.array(x).transpose()]
sig = [[numpy.mean(x_i * x_k) for x_k in x] for x_i in x]
eps = [
numpy.sqrt(numpy.linalg.det(numpy.array(sig)[i : i + 2][:, i : i + 2]))
for i in (0, 2, 4)
]
bet = [sig[i][i] / eps[i / 2] for i in (0, 2, 4)]
alf = [-sig[i][i + 1] / eps[i / 2] for i in (0, 2, 4)]
gam = [sig[i + 1][i + 1] / eps[i / 2] for i in (0, 2, 4)]
def x2twiss_halo(x, sig_cut=(None, None, None)):
"""
Calculate twiss of halo particles (r > sig_cut).
Note halo particles are different for each plane.
Input : x (Nptcl,6), sig_cut
Output: eps,bet,alf,gam
2015.07.30
"""
eps, bet, alf, gam = x2twiss(x)
for k in (0, 2, 4):
x = x.transpose()
r_nom_sq = (
x[k] ** 2 + (alf[k / 2] * x[k] + bet[k / 2] * x[k + 1]) ** 2
) / (bet[k / 2] * eps[k / 2])
x = x.transpose()
x_halo = [x[i] for i in range(len(x)) if r_nom_sq[i] > sig_cut[k / 2] ** 2]
eps_tmp, bet_tmp, alf_tmp, gam_tmp = x2twiss(x_halo)
eps[k / 2] = eps_tmp[k / 2]
bet[k / 2] = bet_tmp[k / 2]
alf[k / 2] = alf_tmp[k / 2]
gam[k / 2] = gam_tmp[k / 2]
return eps, bet, alf, gam
def loss_elem2den(s, loss, file_name_dt="", dlt_dt=5e-6):
"""
This function calculates loss density [W/m] from s and loss
together with the file of the DTL's cell and DT lengths.
L = [s[i] - s[i - 1] for i in range(1, len(s))]
L.insert(0, 0.0)
# Take care DTL part if "file_name_dt" is given
try:
# DTL cell and DT lengths
with open(file_name_dt) as file:
L_cell, L_dt = numpy.array(
[map(float, lin.split()) for lin in file.readlines()]
).transpose()[:2]
dL = list(abs(L_cell - L[i]))
dL_min = min(dL)
if dL_min < dlt_dt:
L[i] = L_dt[dL.index(dL_min)]
Ndt += 1
if Ndt != len(L_dt):
print(
"drift-tube # unmatched, in file: {}, matched: {}".format(
len(L_dt), Ndt
)
)
if loss_den[i] != 0.0 and L[i] == 0.0:
print("Caution: inf loss density at elem # ", i, "!!!!")
if L[k] != 0.0:
loss_den[k] += loss_den[i]
loss_den[i] = 0.0
break
# Calculate loss den
for i in range(len(loss_den)):
if loss_den[i] != 0.0:
loss_den[i] /= L[i]
lngth = [s[i + 1] - s[i] for i in range(len(s) - 1)]
lngth.insert(0, 0.0)
# L=[s[i+1]-s[i] for i in range(len(s)-1)]; L.insert(0,0.0)
# Take care DTL part if "file_name_dt" is given
try:
# Read DTL cell and DT lengths
with open(file_name_dt) as file:
l_cell, l_dt = numpy.array(
[map(float, lin.split()) for lin in file.readlines()]
).transpose()[:2]
for i in range(len(lngth)):
dl = list(abs(l_cell - lngth[i]))
if Ndt != len(l_cell):
print(
"drift-tube # unmatched, in file: {}, matched: {}".format(
len(l_cell), Ndt
)
)
for i in range(len(lngth)):
if lngth[i] == 0.0 and loss[i] != 0.0:
if i == 1:
print("The 1st elem has 0 length and yet has a loss, exiting...")
k += 1
loss[i - k] += loss[i]
loss[i] = 0.0
print("Caution: Inf loss density at elem #", i, "!!!!")
if loss[i] == 0.0:
loss_den.append(0.0)
else:
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
# def loss_elem2den(s,loss,path_file_dt,dlt_dt=5e-6):
# '''
# This function calculates loss density [W/m] from s, loss,
# and the file of the DTL's cell and DT lengths.
# 2015.01.30
# '''
# # Define l
# l=[s[i+1]-s[i] for i in range(len(s)-1)]; l.insert(0,0.0)
# # DTL cell and DT lengths.
# file=open(path_file_dt)
# l_cell,l_dt=array([map(float,lin.split()) for lin in file.readlines()]).transpose()[:2]
# file.close()
# # Replace l's in DTL with DT lengths.
# Ndt=0
# for i in range(len(l)):
# dl=list(abs(l_cell-l[i])); dl_min=min(dl)
# if dl_min<dlt_dt:
# l[i]=l_dt[dl.index(dl_min)]; Ndt+=1
# # Check the replacement.
# if Ndt!=len(l_cell):
# print 'drift-tube # unmatched, in file: '+str(len(l_cell))+', from matching: '+str(Ndt)
# print 'review the threshold, exiting...'; exit()
# # Treat inf loss den.
# for i in range(len(l)):
# if l[i]==0.0 and loss[i]!=0.0:
# if i==1:
# print 'The 1st elem has 0 length and yet has a loss, exiting...'; exit()
# else:
# k=1
# while l[i-k]==0.0: k+=1
# loss[i-k]+=loss[i]; loss[i]=0.0
# print 'Caution: Inf loss density at elem #',i,'!!!!'
# # Finally, convert to loss density.
# loss_den=[]
# for i in range(len(l)):
# if loss[i]==0.0: loss_den.append(0.0)
# else : loss_den.append(loss[i]/l[i])
# return loss_den
def partran_end_all(file_name_in_all, file_name_out):
"""
- Reads multiple partran1.out files and extracts the data at the end.
- The same format as partran1.out allows to be used by PARTRAN class.
- The elem # is replaced with the file #.
2015.11.25
with open(file_name_in_all[0], "r") as file:
header = ""
header += lin
if lin.split()[0] == "####":
break
# Extract output data for each file and write.
with open(file_name_out, "w") as file_out:
file_out.write(header)
for i in range(len(file_name_in_all)):
lin = check_output(["tail", "-1", file_name_in_all[i]]).split()
# Just to maintain the same format as partran1.out. Maybe too much...
for i in range(1, len(lin)):
if "." in lin[i]:
file_out.write(" %13s" % lin[i])
else:
file_out.write(" %10s" % lin[i])
file_out.write("\n")
def _update_field_map_dict(dictionary, folder_path):
for filename in os.listdir(folder_path):
if filename.split(".")[-1] == "edz": # only 1D for now..
key = filename.split(".")[0]
if key not in dictionary: # do not override user selection
dictionary[key] = os.path.join(folder_path, filename)
logging.debug(" Found field map {}: {}".format(key, dictionary[key]))
# ---- Old MATCH and LATTICE classes from the time of IPAC'15 (more complex)
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
# class MATCH:
# '''
# Class for matching. Could be revised.
# 2015.04.26
# '''
# def __init__(self,typ):
# # Instances
# self.typ =typ # Need typ_knob and etc ??
# self.i_diag=[]
# self.i_knob=[]
# # Instances for 'TRAJ'
# self.Rx_inv=None
# self.Ry_inv=None
# class LATTICE:
# '''
# Class to handle a TraceWin lattice file.
# - For now, only for the steerer and BPM. Extend as needed.
# - For now, only for the MEBT. Extend as needed.
# - For now, reading a file for ken. Implement a method for future ??
# 2015.04.12
# '''
# def __init__(self,path_file_lat,path_file_ken):
# # Consts, need to define globally ??
# mass=938.2723; clight=2.99792458
# # Some dict and list (MEBT and DTL are supported for now...)
# dic_class={'QUAD' : QUAD ,
# 'THIN_STEERING' : THIN_STEERING ,
# 'STEERER' : STEERER ,
# 'DIAG_POSITION' : DIAG_POSITION ,
# 'ADJUST' : ADJUST ,
# 'ADJUST_STEERER': ADJUST_STEERER}
# lst_elem =['DRIFT','QUAD','GAP','DTL_CEL',
# 'THIN_STEERING','DIAG_POSITION']
# lst_comm =['ADJUST','ADJUST_STEERER']
# lst_diag =['DIAG_POSITION']
# # Load ken from TW's "Energy(MeV).txt" (vs Nelem, starting from 0)
# with open(path_file_ken,'r') as file:
# ken=[]
# for lin in file:
# try : ken.append(float(lin.split()[2]))
# except: pass
# ken.insert(0,ken[0])
# # Go through the lat file
# with open(path_file_lat) as file:
# lst=[]; i=0; Nelem=0
# for lin in file:
# lin=lin.partition(';')[0] # Remove comment
# if lin.split()!=[]: # Remove empty line
# # Convert an elem/comm to a class (maybe too cryptic...)
# try : lst.append(dic_class[ALLTYPE(lin).typ](lin))
# except: lst.append( ALLTYPE(lin))
# # Assign index
# lst[-1].indx=i; i+=1
# # Assign Nelem
# if lst[-1].typ in lst_elem: Nelem+=1
# lst[-1].Nelem=Nelem
# # Assign gamma, beta, Brho
# gamma=1.0+ken[Nelem]/mass ; lst[-1].gamma=gamma
# beta =sqrt(1.0-1.0/gamma**2) ; lst[-1].beta =beta
# Brho =(10.0/clight)*gamma*beta*mass*1e-3; lst[-1].Brho =Brho
# # End at 'END'
# if lst[-1].typ=='END': break
# # Go through lat again and define matchings, extend/modify as needed
# match={}
# for i in range(len(lst)):
# if lst[i].typ in lst_comm:
# lst,match=lst[i].activate(lst,match)
# # Get index of diagnostics
# for i in range(len(lst)):
# if lst[i].typ in lst_diag:
# try : match[lst[i].fam].i_diag.append(i)
# except: pass
# # Instances
# self.lst =lst
# self.match=match
# def get_tw(self,path_file,flag_no_err=None):
# with open(path_file,'w') as file:
# for lst_i in self.lst:
# if 'ERROR' in lst_i.typ:
# if flag_no_err!=None: pass
# else : file.write(lst_i.get_tw())
# else:
# file.write(lst_i.get_tw())
# ---- Old DENSITY class for the envelope calc only
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
# class DENSITY_ENV:
# '''
# This class is to handle Density_Env.dat.
# - For now, this is intended for an individual file and NOT "tot" file.
# (Something not clear between an individual file and "tot" file...)
# - For now, there are instances of only Nelem, s, x_ave, y_ave.
# (Add others when needed.)
# - For now, tested only for the MEBT.
# - Note there are much more data points than the total number of elems
# due to the slicing of the SC calculation. Use i_elem to extract data
# at the ends of elems, e.g., array(x_ave)[array(i_elem)].
# 2015.03.29
# '''
# def __init__(self,path_file):
# # Go through the file
# Nelem=[]; s=[]; x=[]
# with open(path_file) as file:
# while True:
# try:
# ver=fromfile(file,dtype = numpy.uint16,count=3)[0] # version, year, vlong
# Nelem.append( numpy.fromfile(file,dtype = numpy.uint32,count=2)[1]) # Nrun, Nelem
# s.append( numpy.fromfile(file,dtype = numpy.float32,count=4)[1]) # Ibeam, s, aptx, apty
# numpy.fromfile(file,dtype = numpy.uint32,count=1) # Nslice ??
# x.append(list( numpy.fromfile(file,dtype = numpy.float32,count=7*4)[:2])) # ave & etc
# if ver>=5: numpy.fromfile(file,dtype = numpy.float32,count=7*2) # ver >= 5 part
# if ver>=6: numpy.fromfile(file,dtype = numpy.float32,count=7*2) # ver >= 6 part
# if ver>=7: numpy.fromfile(file,dtype = numpy.float32,count=3*2) # ver >= 7 part
# if ver>=8: numpy.fromfile(file,dtype = numpy.float32,count=1*3) # ver >= 8 part
# numpy.fromfile(file,dtype = numpy.uint64,count=1) # Nptcl
# except:
# break
# # Index of n-th elem end, starting from 0
# i_elem=[0]
# for n in range(1,Nelem[-1]):
# k=1
# while True:
# try : i_elem.append(Nelem.index(n+k)-1); break
# except: k+=1
# i_elem.append(len(Nelem)-1)
# # Instances
# self.Nelem=Nelem
# self.s =s
# self.x =list(1e3*array(x).transpose()[0])
# self.y =list(1e3*array(x).transpose()[1])
# # Additional instances
# self.i_elem=i_elem
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
# def clean_lat(path_file):
# '''
# Remove comments, names, and etc from a TraceWin file.
# To convert the output to str, use: [' '.join(l)+'\n' for l in lat].
# The current version not capitalizing.
# Output: Lattice in a list
# 2015.03.18
# '''
# with open(path_file) as file:
# lat=[]
# for lin in file:
# lin=lin.partition(';')[0] # Remove comment
# #lin=lin.upper().partition(';')[0] # Remove comment and capitalize
# if ':' in lin: lin=lin.partition(':')[-1].split() # Remove name
# else : lin=lin.split()
# if lin!=[]: # Avoid empty line
# lat.append(lin)
# if lin[0].upper()=='END': break
# return lat
# def get_steerer(path_file):
# '''
# Extract steerer values from "Adjusted_Values.txt" of TraceWin.
# Note the key is the "lat elem #" where commands are not counted.
# Output is dic of {lat elem #:[val]} or {lat elem #:[val_x,val_y]}
# 2014.10.17
# '''
# file=open(path_file)
# i_elem=0; dic_steerer_val={}
# for lin in file.readlines():
# lin=lin.split()
# for i in range(len(lin)):
# if lin[i]==':':
# if int(lin[i-1])!=i_elem:
# i_elem=int(lin[i-1])
# dic_steerer_val[i_elem]=[lin[i+1]]
# else:
# dic_steerer_val[i_elem].append(lin[i+1])
# file.close()
# return dic_steerer_val
# def apply_steerer(lat,dic_steerer_val,lst_typ_elem):
# '''
# Apply steerer values to a TraceWin lattice.
# Input: lat from "clean_lat, steerer dic from "get_steerer", list of lat elem types
# 2014.10.17
# '''
# #--
# i_elem=0; dic_steerer={}
# for i in range(len(lat)):
# if lat[i][0] in lst_typ_elem:
# i_elem+=1
# if lat[i][0]=='STEERER':
# if i_elem+1 in dic_steerer_val.keys():
# if len(dic_steerer_val[i_elem+1])==2:
# dic_steerer[i]=[[1,dic_steerer_val[i_elem+1][0]],[2,dic_steerer_val[i_elem+1][1]]]
# else:
# for k in range(i)[::-1]:
# if lat[k][0]=='ADJUST_STEERER_BX':
# dic_steerer[i]=[[1,dic_steerer_val[i_elem+1][0]]]; break
# if lat[k][0]=='ADJUST_STEERER_BY':
# dic_steerer[i]=[[2,dic_steerer_val[i_elem+1][0]]]; break
# if lat[i][0]=='THIN_STEERING':
# if i_elem in dic_steerer_val.keys():
# if len(dic_steerer_val[i_elem])==2:
# dic_steerer[i]=[[1,dic_steerer_val[i_elem][0]],[2,dic_steerer_val[i_elem][1]]]
# else:
# for k in range(i)[::-1]:
# if lat[k][0]=='ADJUST' and lat[k][2]=='1':
# dic_steerer[i]=[[1,dic_steerer_val[i_elem][0]]]; break
# if lat[k][0]=='ADJUST' and lat[k][2]=='2':
# dic_steerer[i]=[[2,dic_steerer_val[i_elem][0]]]; break
# for i in sorted(dic_steerer.keys()):
# for k in range(len(dic_steerer[i])):
# lat[i][dic_steerer[i][k][0]]=dic_steerer[i][k][1]
# return lat
# def apply_multipole(lat,typ,r,b3=0.0,b4=0.0,b5=0.0,b6=0.0):
# '''
# Apply multipole components to quads in a TraceWin lattice file.
# Note all the quads in all the sections are affected in the same way.
# Input : lat from "clean_lat, reference r, b3-b6, and dist type
# ("fix", "uniform", "gauss")
# 2015.01.08
# '''
# bn=(b3,b4,b5,b6)
# for i in range(len(lat)):
# if lat[i][0]=='QUAD':
# # Adjust format
# if len(lat[i])<5:
# for k in range(5): lat[i].append('0')
# else:
# lat[i]=lat[i][:5]
# for k in range(4): lat[i].append('0')
# # Fixed value
# if typ=='fix':
# for k in range(4):
# lat[i][k+5]=str( bn[k]/r**(k+1)*float(lat[i][2]))
# # Uniform dist
# if typ=='uniform':
# for k in range(4):
# lat[i][k+5]=str((2.0*rand()-1.0)*bn[k]/r**(k+1)*float(lat[i][2]))
# # Gaussian dist
# if typ=='gauss':
# for k in range(4):
# lat[i][k+5]=str( randn()*bn[k]/r**(k+1)*float(lat[i][2]))
# return lat
# Extract output Twiss and halo from partran1.out
# for lin in file.readlines():
# lin=lin.split()
# if lin[0]=='####':
# idx_gamma=lin.index('gama-1')
# idx_eps =[lin.index(p) for p in ("ex" ,"ey" ,"ezdp" )]
# idx_bet =[lin.index(p) for p in ("SizeX","SizeY","SizeZ")]
# idx_alf =[lin.index(p) for p in ("sxx'" ,"syy'" ,"szdp" )]
# idx_hal =[lin.index(p) for p in ("hx" ,"hy" ,"hp" )]
# file.close()
# gamma=float(lin[idx_gamma])+1.0
# eps =array([float(lin[idx]) for idx in idx_eps])
# bet =array([float(lin[idx])**2 for idx in idx_bet]) # Actually Sig_{i,i}
# alf =array([float(lin[idx]) for idx in idx_alf]) # Actually Sig_{i,i+1}
# hal =array([float(lin[idx]) for idx in idx_hal])
# bet= bet/eps*sqrt(gamma**2-1.0); bet[2]*=gamma**2 # z-dp => z-z'
# alf=-alf/eps*sqrt(gamma**2-1.0)
# gam= (1.0+alf**2)/bet
# class PROJECT_SING_CORE:
# '''
# - Use this for simultaneous mult 1-core-simulations (w/ tracelx64).
# - Initial parameters are not meant to be changed, unlike "PROJECT".
# (Each run should have its own project.)
# - Make sure the project/lattice file is in the calc dir.
# - Changing the calc dir option not working for tracelx64.
# def __init__(self,path_cal,
# file_name_proj='proj.ini',file_name_lat='lat.dat',seed=None,flag_hide=None):
# # Make sure "path_cal" ends with "/".
# if path_cal[-1]!='/': path_cal+='/'
# # Check the calc dir exists.
# if isdir(path_cal)==False: print 'Calc dir does not exist !!!! Exiting ...'; exit()
# # Check "tracelx64" is in the calc dir.
# if isfile(path_cal+'tracelx64')==False:
# try : system('cp /opt/cea/tracelx64 '+path_cal)
# except: print 'tracelx64 not in calc dir !!!! Exiting ...'; exit()
# # Check "tracewin.key" is in the calc dir.
# if isfile(path_cal+'tracewin.key')==False:
# try : system('cp /opt/cea/tracewin.key '+path_cal)
# except: print 'tracewin.key not in calc dir !!!! Exiting ...'; exit()
# # Check the project/lattice file is in the calc dir.
# if '/' not in file_name_proj: file_name_proj=path_cal+file_name_proj
# if '/' not in file_name_lat : file_name_lat =path_cal+file_name_lat
# if isfile(file_name_proj)==False: print 'project file not in calc dir !!!! Exiting ...'; exit()
# if isfile(file_name_lat) ==False: print 'lattice file not in calc dir !!!! Exiting ...'; exit()
# # Instances (Add more options as needed.)
# self.path_cal =path_cal
# self.file_name_proj=file_name_proj
# self.file_name_lat =file_name_lat
# self.seed =seed
# self.flag_hide =flag_hide
# opt_exe =self.path_cal+'tracelx64 '+self.file_name_proj
# if self.file_name_lat!=None: opt_exe+=' dat_file=' +self.file_name_lat
# if self.seed !=None: opt_exe+=' random_seed='+self.seed
# if self.flag_hide !=None: opt_exe+=' hide'