diff --git a/ess/lib_tw.py b/ess/lib_tw.py
index 2df3892852dcc072cf7faf09ea346a889443c63f..c5d20a3c719bae29be8df7e7b607e8bee78fb136 100644
--- a/ess/lib_tw.py
+++ b/ess/lib_tw.py
@@ -13,7 +13,7 @@ from __future__ import print_function
import numpy
-from struct import pack
+from struct import pack, unpack
from itertools import chain
from subprocess import check_output
@@ -204,16 +204,16 @@ class LATTICE:
if self.lst[k].apt is not None:
self.lst[i].apt = self.lst[k].apt
break
- except:
+ except IndexError:
pass
else:
for k in range(i)[::-1]:
try:
self.lst[i].apt = self.lst[k].apt
break
- except:
+ except IndexError:
pass
- except:
+ except IndexError:
pass
def update_gamma(self):
@@ -225,7 +225,7 @@ class LATTICE:
self.lst[i].gamma = self.lst[i - 1].gamma
try:
self.lst[i].update_gamma()
- except:
+ except IndexError:
pass
def update_st(self, file_name):
@@ -278,7 +278,7 @@ class LATTICE:
if i == "ERROR" and lin[0] == "BEAM":
continue
settings = {}
- for j in xrange(len(lin[2:]) / 3):
+ for j in range(len(lin[2:]) / 3):
k = int(lin[2 + 3 * j])
val = float(lin[4 + 3 * j])
if k in settings:
@@ -414,13 +414,13 @@ class PROJECT:
def exe(self):
opt_exe = "TraceWin64_noX11 " + self.file_name_proj
- if self.file_name_lat != None:
+ if self.file_name_lat is not None:
opt_exe += " dat_file=" + self.file_name_lat
- if self.path_cal != None:
+ if self.path_cal is not None:
opt_exe += " path_cal=" + self.path_cal
- if self.seed != None:
+ if self.seed is not None:
opt_exe += " random_seed=" + self.seed
- if self.flag_hide != None:
+ if self.flag_hide is not None:
opt_exe += " hide"
# if self.path_cal!=None:
@@ -435,20 +435,19 @@ class PROJECT:
class PARTRAN:
"""
Note:
-
+
- The list not complete. Add parameters as needed.
-
+
History:
-
+
- 2016.02.17: Changed how to identify the line of indices.
- 2016.02.17: Added a logic to avoid #/0 for LEBT.
-
+
"""
def __init__(self, file_name):
# Consts to convert phs to z.
- mass = 938.272029
c = 2.99792458
freq = 352.21
@@ -515,13 +514,13 @@ class PARTRAN:
# Way around for 0 emitt
for i in range(len(self.s)):
if self.epsx[i] == 0.0:
- self.epsx[i] = inf
+ self.epsx[i] = numpy.inf
if self.epsy[i] == 0.0:
- self.epsy[i] = inf
+ self.epsy[i] = numpy.inf
if self.epsz[i] == 0.0:
- self.epsz[i] = inf
+ self.epsz[i] = numpy.inf
if self.epsp[i] == 0.0:
- self.epsp[i] = inf
+ self.epsp[i] = numpy.inf
# Additional instances
self.gamma = data[idx_gamma] + 1.0
@@ -538,13 +537,13 @@ class PARTRAN:
# Set inf emitt back to 0
for i in range(len(self.s)):
- if self.epsx[i] == inf:
+ if self.epsx[i] == numpy.inf:
self.epsx[i] = 0.0
- if self.epsy[i] == inf:
+ if self.epsy[i] == numpy.inf:
self.epsy[i] = 0.0
- if self.epsz[i] == inf:
+ if self.epsz[i] == numpy.inf:
self.epsz[i] = 0.0
- if self.epsp[i] == inf:
+ if self.epsp[i] == numpy.inf:
self.epsp[i] = 0.0
# Convert to list (not necessary?)
@@ -602,11 +601,11 @@ class DST:
# Read the file
with open(file_name) as file:
- dummy = numpy.fromfile(file, dtype=numpy.uint8, count=2)
+ numpy.fromfile(file, dtype=numpy.uint8, count=2)
Nptcl = numpy.fromfile(file, dtype=numpy.uint32, count=1)[0]
Ibeam = numpy.fromfile(file, dtype=numpy.float64, count=1)[0]
freq = numpy.fromfile(file, dtype=numpy.float64, count=1)[0]
- dummy = numpy.fromfile(file, dtype=numpy.uint8, count=1)
+ numpy.fromfile(file, dtype=numpy.uint8, count=1)
x = (
numpy.fromfile(file, dtype=numpy.float64, count=Nptcl * 6)
.reshape(Nptcl, 6)
@@ -656,11 +655,11 @@ class DENSITY:
Nelem : loss_num_(..), loss_pow_(..)
Nelem x Nstep: den
-
+
- History:
-
+
* 2015.11.12: Baseline ver
- * 2016.03.29: Adapted to ver 9 (apt includes shifts)
+ * 2016.03.29: Adapted to ver 9 (apt includes shifts)
"""
def __init__(self, file_name):
@@ -702,32 +701,48 @@ class DENSITY:
while True:
try:
# Partran and envelope
- ver, year, flag_long = fromfile(file, dtype=uint16, count=3)
- Nrun = fromfile(file, dtype=uint32, count=1)[0]
- idx_elem.append(fromfile(file, dtype=uint32, count=1)[0])
- Ibeam.append(fromfile(file, dtype=float32, count=1)[0])
- s.append(fromfile(file, dtype=float32, count=1)[0])
+ ver, year, flag_long = numpy.fromfile(
+ file, dtype=numpy.uint16, count=3
+ )
+ Nrun = numpy.fromfile(file, dtype=numpy.uint32, count=1)[0]
+ idx_elem.append(
+ numpy.fromfile(file, dtype=numpy.uint32, count=1)[0]
+ )
+ Ibeam.append(numpy.fromfile(file, dtype=numpy.float32, count=1)[0])
+ s.append(numpy.fromfile(file, dtype=numpy.float32, count=1)[0])
if ver >= 9:
- apt.append(fromfile(file, dtype=float32, count=4))
+ apt.append(numpy.fromfile(file, dtype=numpy.float32, count=4))
else:
- apt.append(fromfile(file, dtype=float32, count=2))
- Nstep = fromfile(file, dtype=uint32, count=1)[0]
- cent_ave.append(fromfile(file, dtype=float32, count=7))
- cent_rms.append(fromfile(file, dtype=float32, count=7))
- xmax.append(fromfile(file, dtype=float32, count=7))
- xmin.append(fromfile(file, dtype=float32, count=7))
+ apt.append(numpy.fromfile(file, dtype=numpy.float32, count=2))
+ Nstep = numpy.fromfile(file, dtype=numpy.uint32, count=1)[0]
+ cent_ave.append(numpy.fromfile(file, dtype=numpy.float32, count=7))
+ cent_rms.append(numpy.fromfile(file, dtype=numpy.float32, count=7))
+ xmax.append(numpy.fromfile(file, dtype=numpy.float32, count=7))
+ xmin.append(numpy.fromfile(file, dtype=numpy.float32, count=7))
if ver > 5:
- sig_ave.append(fromfile(file, dtype=float32, count=7))
- sig_rms.append(fromfile(file, dtype=float32, count=7))
+ sig_ave.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=7)
+ )
+ sig_rms.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=7)
+ )
if ver >= 6:
- cent_min.append(fromfile(file, dtype=float32, count=7))
- cent_max.append(fromfile(file, dtype=float32, count=7))
+ cent_min.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=7)
+ )
+ cent_max.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=7)
+ )
if ver >= 7:
- eps_ave.append(fromfile(file, dtype=float32, count=3))
- eps_rms.append(fromfile(file, dtype=float32, count=3))
+ eps_ave.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=3)
+ )
+ eps_rms.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=3)
+ )
if ver >= 8:
- accpt.append(fromfile(file, dtype=float32, count=3))
- Nptcl.append(fromfile(file, dtype=uint64, count=1)[0])
+ accpt.append(numpy.fromfile(file, dtype=numpy.float32, count=3))
+ Nptcl.append(numpy.fromfile(file, dtype=numpy.uint64, count=1)[0])
# Partran only
if Nptcl[-1] > 0:
loss_num.append([])
@@ -736,56 +751,74 @@ class DENSITY:
den_pow.append([])
for n in range(Nrun):
loss_num[-1].append(
- fromfile(file, dtype=uint64, count=1)[0]
+ numpy.fromfile(file, dtype=numpy.uint64, count=1)[0]
)
loss_pow[-1].append(
- fromfile(file, dtype=float32, count=1)[0]
+ numpy.fromfile(file, dtype=numpy.float32, count=1)[0]
)
loss_num_ave.append(sum(loss_num[-1]))
- loss_num_rms.append(fromfile(file, dtype=uint64, count=1)[0])
- loss_num_min.append(fromfile(file, dtype=uint64, count=1)[0])
- loss_num_max.append(fromfile(file, dtype=uint64, count=1)[0])
+ loss_num_rms.append(
+ numpy.fromfile(file, dtype=numpy.uint64, count=1)[0]
+ )
+ loss_num_min.append(
+ numpy.fromfile(file, dtype=numpy.uint64, count=1)[0]
+ )
+ loss_num_max.append(
+ numpy.fromfile(file, dtype=numpy.uint64, count=1)[0]
+ )
loss_pow_ave.append(sum(loss_pow[-1]))
- loss_pow_rms.append(fromfile(file, dtype=float64, count=1)[0])
- loss_pow_min.append(fromfile(file, dtype=float32, count=1)[0])
- loss_pow_max.append(fromfile(file, dtype=float32, count=1)[0])
+ loss_pow_rms.append(
+ numpy.fromfile(file, dtype=numpy.float64, count=1)[0]
+ )
+ loss_pow_min.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=1)[0]
+ )
+ loss_pow_max.append(
+ numpy.fromfile(file, dtype=numpy.float32, count=1)[0]
+ )
for k in range(7):
if flag_long == 1:
den[-1].append(
- fromfile(file, dtype=uint64, count=Nstep)
+ numpy.fromfile(
+ file, dtype=numpy.uint64, count=Nstep
+ )
)
else:
den[-1].append(
- fromfile(file, dtype=uint32, count=Nstep)
+ numpy.fromfile(
+ file, dtype=numpy.uint32, count=Nstep
+ )
)
if Ibeam[-1] > 0:
for k in range(3):
den_pow[-1].append(
- fromfile(file, dtype=float32, count=Nstep)
+ numpy.fromfile(
+ file, dtype=numpy.float32, count=Nstep
+ )
)
# print Nrun,Nptcl[-1],idx_elem[-1] # Diag
- except:
+ except IndexError:
break
# -- Reshape arrays
- apt = swapaxes(apt, 1, 0)
- accpt = swapaxes(accpt, 1, 0)
- cent_ave = swapaxes(cent_ave, 1, 0)
- cent_rms = swapaxes(cent_rms, 1, 0)
- cent_max = swapaxes(cent_max, 1, 0)
- cent_min = swapaxes(cent_min, 1, 0)
- sig_ave = swapaxes(sig_ave, 1, 0)
- sig_rms = swapaxes(sig_rms, 1, 0)
- eps_ave = swapaxes(eps_ave, 1, 0)
- eps_rms = swapaxes(eps_rms, 1, 0)
- xmax = swapaxes(xmax, 1, 0)
- xmin = swapaxes(xmin, 1, 0)
+ apt = numpy.swapaxes(apt, 1, 0)
+ accpt = numpy.swapaxes(accpt, 1, 0)
+ cent_ave = numpy.swapaxes(cent_ave, 1, 0)
+ cent_rms = numpy.swapaxes(cent_rms, 1, 0)
+ cent_max = numpy.swapaxes(cent_max, 1, 0)
+ cent_min = numpy.swapaxes(cent_min, 1, 0)
+ sig_ave = numpy.swapaxes(sig_ave, 1, 0)
+ sig_rms = numpy.swapaxes(sig_rms, 1, 0)
+ eps_ave = numpy.swapaxes(eps_ave, 1, 0)
+ eps_rms = numpy.swapaxes(eps_rms, 1, 0)
+ xmax = numpy.swapaxes(xmax, 1, 0)
+ xmin = numpy.swapaxes(xmin, 1, 0)
if Nptcl[0] > 0:
- loss_num = swapaxes(loss_num, 1, 0)
- loss_pow = swapaxes(loss_pow, 1, 0)
- den = swapaxes(den, 1, 0)
- den_pow = swapaxes(den_pow, 1, 0)
+ loss_num = numpy.swapaxes(loss_num, 1, 0)
+ loss_pow = numpy.swapaxes(loss_pow, 1, 0)
+ den = numpy.swapaxes(den, 1, 0)
+ den_pow = numpy.swapaxes(den_pow, 1, 0)
# -- Take care ave and rms
@@ -819,21 +852,21 @@ class DENSITY:
# -- Define std (around to avoid sqrt(-eps))
if Nrun > 1:
- cent_std = numpy.sqrt(around(cent_rms ** 2 - cent_ave ** 2, 12))
- sig_std = numpy.sqrt(around(sig_rms ** 2 - sig_ave ** 2, 12))
- eps_std = numpy.sqrt(around(eps_rms ** 2 - eps_ave ** 2, 12))
- cent_std = nan_to_num(cent_std) # Replace nan with 0
- sig_std = nan_to_num(sig_std) # Replace nan with 0
- eps_std = nan_to_num(eps_std) # Replace nan with 0
+ cent_std = numpy.sqrt(numpy.around(cent_rms ** 2 - cent_ave ** 2, 12))
+ sig_std = numpy.sqrt(numpy.around(sig_rms ** 2 - sig_ave ** 2, 12))
+ eps_std = numpy.sqrt(numpy.around(eps_rms ** 2 - eps_ave ** 2, 12))
+ cent_std = numpy.nan_to_num(cent_std) # Replace nan with 0
+ sig_std = numpy.nan_to_num(sig_std) # Replace nan with 0
+ eps_std = numpy.nan_to_num(eps_std) # Replace nan with 0
if Nptcl[0] > 0:
loss_num_std = numpy.sqrt(
- around(loss_num_rms ** 2 - loss_num_ave ** 2, 16)
+ numpy.around(loss_num_rms ** 2 - loss_num_ave ** 2, 16)
)
loss_pow_std = numpy.sqrt(
- around(loss_pow_rms ** 2 - loss_pow_ave ** 2, 16)
+ numpy.around(loss_pow_rms ** 2 - loss_pow_ave ** 2, 16)
)
- loss_num_std = nan_to_num(loss_num_std) # Replace nan with 0
- loss_pow_std = nan_to_num(loss_pow_std) # Replace nan with 0
+ loss_num_std = numpy.nan_to_num(loss_num_std) # Replace nan with 0
+ loss_pow_std = numpy.nan_to_num(loss_pow_std) # Replace nan with 0
# -- Convert to list (just in case...)
@@ -1000,7 +1033,7 @@ def plt2x(path_file):
x.append(
[[unpack("f", file.read(4))[0] for l in range(7)] for k in range(Nptcl)]
)
- except:
+ except IndexError:
break
file.close()
@@ -1054,8 +1087,8 @@ def x2twiss(x):
2015.07.30
"""
- x = [x_i - mean(x_i) for x_i in numpy.array(x).transpose()]
- sig = [[mean(x_i * x_k) for x_k in x] for x_i in x]
+ 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)
@@ -1130,7 +1163,7 @@ def loss_elem2den(s, loss, file_name_dt="", dlt_dt=5e-6):
len(L_dt), Ndt
)
)
- except:
+ except FileNotFoundError:
pass
# Take care loss in 0 length elem
@@ -1154,8 +1187,8 @@ def loss_elem2den(s, loss, file_name_dt="", dlt_dt=5e-6):
sys.exit()
# Define L
- l = [s[i + 1] - s[i] for i in range(len(s) - 1)]
- l.insert(0, 0.0)
+ 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
@@ -1167,11 +1200,11 @@ def loss_elem2den(s, loss, file_name_dt="", dlt_dt=5e-6):
).transpose()[:2]
# Replace cell lengths with DT lengths
Ndt = 0
- for i in range(len(l)):
- dl = list(abs(l_cell - l[i]))
+ for i in range(len(lngth)):
+ dl = list(abs(l_cell - lngth[i]))
dl_min = min(dl)
if dl_min < dlt_dt:
- l[i] = l_dt[dl.index(dl_min)]
+ lngth[i] = l_dt[dl.index(dl_min)]
Ndt += 1
# Check
if Ndt != len(l_cell):
@@ -1180,18 +1213,18 @@ def loss_elem2den(s, loss, file_name_dt="", dlt_dt=5e-6):
len(l_cell), Ndt
)
)
- except:
+ except FileNotFoundError:
pass
# Take care inf loss den
- for i in range(len(l)):
- if l[i] == 0.0 and loss[i] != 0.0:
+ 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...")
sys.exit()
else:
k = 1
- while l[i - k] == 0.0:
+ while lngth[i - k] == 0.0:
k += 1
loss[i - k] += loss[i]
loss[i] = 0.0
@@ -1199,11 +1232,11 @@ def loss_elem2den(s, loss, file_name_dt="", dlt_dt=5e-6):
# Finally, convert to loss density
loss_den = []
- for i in range(len(l)):
+ for i in range(len(lngth)):
if loss[i] == 0.0:
loss_den.append(0.0)
else:
- loss_den.append(loss[i] / l[i])
+ loss_den.append(loss[i] / lngth[i])
return loss_den
@@ -1427,16 +1460,16 @@ def _update_field_map_dict(dictionary, folder_path):
# with open(path_file) as file:
# while True:
# try:
-# ver=fromfile(file,dtype=uint16,count=3)[0] # version, year, vlong
-# Nelem.append(fromfile(file,dtype=uint32,count=2)[1]) # Nrun, Nelem
-# s.append(fromfile(file,dtype=float32,count=4)[1]) # Ibeam, s, aptx, apty
-# fromfile(file,dtype=uint32,count=1) # Nslice ??
-# x.append(list(fromfile(file,dtype=float32,count=7*4)[:2])) # ave & etc
-# if ver>=5: fromfile(file,dtype=float32,count=7*2) # ver >= 5 part
-# if ver>=6: fromfile(file,dtype=float32,count=7*2) # ver >= 6 part
-# if ver>=7: fromfile(file,dtype=float32,count=3*2) # ver >= 7 part
-# if ver>=8: fromfile(file,dtype=float32,count=1*3) # ver >= 8 part
-# fromfile(file,dtype=uint64,count=1) # Nptcl
+# 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