def init_machine(latfile):
"""Initialize FLAME machine.
Parameters
----------
latfile :
FLAME lattice file.
Returns
-------
tuple
Tuple of ``(m, s)``, where ``m`` is FLAME machine instance,
and ``s`` is initial machine states.
"""
try:
with open(latfile, 'rb') as f:
m = Machine(f)
s = m.allocState({})
m.propagate(s, 0, 1)
_LOGGER.info("ModelFlame: Initialization succeeded.")
return m, s
except Exception as err:
_LOGGER.error(err)
_LOGGER.warning(
"ModelFlame: Lattice file is not valid, do it manually.")
return None, None
def test_init_with_mix(self):
with open(self.latfile, 'rb') as f:
m = Machine(f)
ms = BeamState(machine=m, latfile=self.latfile)
s = m.allocState({})
m.propagate(s, 0, 1)
compare_mstates(self, ms, s)
def test_generate_latfile_original2(self):
# TEST LATFILE
fout1_file = generate_latfile(self.m, latfile=self.latfile1)
with open(fout1_file, 'rb') as f:
f_str = f.read().strip().decode()
self.assertEqual(f_str, self.fout1_str)
with open(fout1_file, 'rb') as f:
m = Machine(f)
s = m.allocState({})
r = m.propagate(s, 0, len(m), range(len(m)))
# TEST RESULTS
for i in range(len(m)):
s1, s0 = r[i][1], self.r[i][1]
for k in self.ref_keys:
self.assertEqual(getattr(s1, k), getattr(s0, k))
for k in self.keys:
self.assertEqual(
getattr(s1, k).tolist(),
getattr(s0, k).tolist())
for k in self.keys_ps:
self.assertEqual(
getattr(s1, k).tolist(),
getattr(s0, k).tolist())
class TestGetIndexByType(unittest.TestCase):
def setUp(self):
latfile = os.path.join(curdir, 'lattice/test_0.lat')
self.latfile = make_latfile(latfile)
with open(self.latfile, 'rb') as f:
self.m = Machine(f)
def test_wrong_type(self):
etyp = ''
e = get_index_by_type(type=etyp, latfile=self.latfile)
self.assertEqual(e, {etyp: []})
etyp1 = 'no_exist_type'
e1 = get_index_by_type(type=etyp1, latfile=self.latfile)
self.assertEqual(e1, {etyp1: []})
def test_one_type(self):
for etyp in get_all_types(latfile=self.latfile):
e = get_index_by_type(type=etyp, latfile=self.latfile)
self.assertEqual(e, {etyp: self.m.find(type=etyp)})
def test_multi_types(self):
all_types = get_all_types(latfile=self.latfile)
for n in range(2, len(all_types)):
etyps = [random.choice(all_types) for _ in range(n)]
e = get_index_by_type(type=etyps, latfile=self.latfile)
e0 = {t: self.m.find(type=t) for t in etyps}
self.assertEqual(e, e0)
class TestGetIndexByName(unittest.TestCase):
def setUp(self):
latfile = os.path.join(curdir, 'lattice/test_0.lat')
self.latfile = make_latfile(latfile)
with open(self.latfile, 'rb') as f:
self.m = Machine(f)
def test_wrong_name(self):
ename = ''
e = get_index_by_name(name=ename, latfile=self.latfile)
self.assertEqual(e, {ename: []})
ename1 = 'no_exist_name'
e = get_index_by_name(name=ename1, latfile=self.latfile)
self.assertEqual(e, {ename1: []})
def test_one_name(self):
""" test_one_name: repeat for 10 times, each with random name
"""
all_names = get_all_names(latfile=self.latfile)
for n in range(10):
ename = random.choice(all_names)
e = get_index_by_name(name=ename, latfile=self.latfile)
self.assertEqual(e, {ename: self.m.find(name=ename)})
def test_multi_names(self):
""" test_multi_names: test names list length of 2~10
"""
all_names = get_all_names(latfile=self.latfile)
for n in range(2, 10):
enames = [random.choice(all_names) for _ in range(n)]
e = get_index_by_name(name=enames, latfile=self.latfile)
e0 = {n: self.m.find(name=n) for n in enames}
self.assertEqual(e, e0)
def test_set_bmstate(self):
fm_none = ModelFlame()
self.assertIsNone(fm_none.bmstate)
with open(self.testfile, 'rb') as f:
m = Machine(f)
s = m.allocState({})
m.propagate(s, 0, 1)
fm_none.bmstate = s
compare_mstates(self, fm_none.bmstate, s)
def test_transmat(self):
with open(self.latfile, 'rb') as f:
m = Machine(f)
s = m.allocState({})
m.propagate(s, 0, 10)
ms = BeamState(s)
left_val = ms.transfer_matrix
right_val = s.transmat
self.assertTrue((left_val == right_val).all())
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
self.refIonZ = S.ref_IonZ
self.IonZ = S.IonZ # list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = np.zeros((len(self.M), 9))
self.LD2 = np.zeros((len(self.M), 7))
# store element position data
self.bpmls = []
self.corls = []
self.cavls = []
self.solls = []
self.cavpara = []
self.solpara = []
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm': self.bpmls.append(i)
elif elem['type'] == 'orbtrim': self.corls.append(i)
elif elem['type'] == 'rfcavity':
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'], elem['phi']])
elif elem['type'] == 'solenoid':
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info', 'w') as f:
iteminfo = [
len(self.M), self.bpmls, self.corls, self.cavls, self.solls
]
cPickle.dump(iteminfo, f)
def test_init_with_s2(self):
""" test_init_with_s2: s is None
"""
with open(self.latfile, 'rb') as f:
m = Machine(f)
s = m.allocState({})
m.propagate(s, 0, 1)
ms = BeamState()
ms.state = s
compare_mstates(self, ms, s)
def test_generate_latfile_update1(self):
idx = 80
self.m.reconfigure(idx, {'theta_x': 0.1})
fout2_file = generate_latfile(self.m, latfile=self.latfile2)
with open(fout2_file) as f:
self.assertEqual(f.read().strip(), self.fout2_str)
with open(fout2_file, 'rb') as f:
m = Machine(f)
self.assertEqual(m.conf(idx)['theta_x'], 0.1)
def test_generate_latfile_update2(self):
idx = 0
s = self.m.allocState({})
self.m.propagate(s, 0, 1)
s.moment0 = [[0.1], [0.1], [0.1], [0.1], [0.1], [0.1], [1.0]]
fout2_file = generate_latfile(self.m, latfile=self.latfile2, state=s)
with open(fout2_file, 'rb') as f:
m = Machine(f)
assertAEqual(m.conf()['P0'], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0])
def test_run_3(self):
""" test run_3: test initial states
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
m0.propagate(s0, 0, 1)
fm = ModelFlame(latfile)
r, s = fm.run(from_element=0, to_element=0)
compare_mstates(self, s0, s)
def test_run_1(self):
""" test_run_1: propagate from the first to last, monitor None
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
m0.propagate(s0, 0, len(m0))
fm = ModelFlame(latfile)
r,s = fm.run()
self.assertEqual(r, [])
compare_mstates(self, s, s0)
def parseFrom(self, lattice_file_name):
with open(lattice_file_name, 'rb') as f:
self.machine = Machine(f)
lattice_dict = self.machine.conf()
self.lattice = copy.deepcopy(lattice_dict['elements'])
self.linename = lattice_dict['name']
'''['AMU',
'Eng_Data_Dir',
'HdipoleFitMode',
'IonChargeStates',
'IonEk',
'IonEs',
'IonW',
'IonZ',
'MpoleLevel',
'NCharge',
'Stripper_IonChargeStates',
'Stripper_NCharge',
'sim_type']
'''
self.AMU = lattice_dict['AMU']
self.Eng_Data_Dir = lattice_dict['Eng_Data_Dir']
self.HdipoleFitMode = float(lattice_dict.get("HdipoleFitMode", 1.0))
self.IonChargeStates = lattice_dict['IonChargeStates']
self.IonEk = lattice_dict['IonEk']
self.IonEs = lattice_dict['IonEs']
self.IonW = lattice_dict.get('IonW', -1)
self.IonZ = lattice_dict.get('IonZ', -1)
self.MpoleLevel = float(lattice_dict.get("MpoleLevel", 2.0))
self.NCharge = lattice_dict['NCharge']
self.Stripper_IonChargeStates = lattice_dict[
'Stripper_IonChargeStates']
self.Stripper_NCharge = lattice_dict['Stripper_NCharge']
self.sim_type = lattice_dict['sim_type']
n_charge_state = len(self.IonChargeStates)
self.BaryCenter = []
self.SMatrix = []
for i in range(n_charge_state):
temp = 'BaryCenter{}'.format(i)
self.BaryCenter.append(lattice_dict[temp])
temp = 'S{}'.format(i)
self.SMatrix.append(lattice_dict[temp])
for ele in self.lattice:
self.addElement(ele['name'], ele['type'], **ele)
self.appendToBeamline(self.linename, ele['name'].upper())
self.setUseLine(self.linename)
def _create_flame_machine(self, lat_file):
try:
self._machine = Machine(open(lat_file, 'r'))
except IOError as e:
print("Failed to open {fn}: {info}".format(fn=e.filename,
info=e.args[-1]))
sys.exit(1)
except (RuntimeError, KeyError) as e:
print("Cannot parse lattice, " + e.args[-1])
sys.exit(1)
except:
print("Failed to create machine")
sys.exit(1)
def test_run_6(self):
""" test_run_6: optional monitor setting 'all'
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
r0 = m0.propagate(s0, 0, len(m0), observe=range(len(m0)))
r,s = fm.run(monitor='all')
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def setUp(self):
testfile = os.path.join(curdir, 'lattice/test_0.lat')
out1file = os.path.join(curdir, 'lattice/out1_0.lat')
out2file = os.path.join(curdir, 'lattice/out2_0.lat')
self.testfile = make_latfile(testfile)
self.out1file = make_latfile(out1file)
self.out2file = make_latfile(out2file)
with open(self.testfile, 'rb') as f:
self.m = Machine(f)
s0 = self.m.allocState({})
self.r = self.m.propagate(s0, 0, len(self.m), range(len(self.m)))
self.s = s0
k = [
'beta', 'bg', 'gamma', 'IonEk', 'IonEs', 'IonQ', 'IonW', 'IonZ',
'phis', 'SampleIonK'
]
self.keys = [i for i in k]
self.ref_keys = ['ref_{}'.format(i) for i in k] + ['pos']
self.keys_ps = ['moment0', 'moment0_env', 'moment0_rms', 'moment1']
with open(self.out1file, 'rb') as f:
self.fout1_str = f.read().strip().decode()
with open(self.out2file, 'rb') as f:
self.fout2_str = f.read().strip().decode()
self.latfile0 = os.path.join(curdir, 'lattice/test_1.lat')
self.latfile1 = os.path.join(curdir, 'lattice/out1_1.lat')
self.latfile2 = os.path.join(curdir, 'lattice/out2_1.lat')
self.testcfile = os.path.join(curdir, 'lattice/test_c.lat')
self.out3file = os.path.join(curdir, 'lattice/out3.lat')
with open(self.out3file, 'rb') as f:
self.fout3_str = f.read().strip().decode()
with open(self.testcfile, 'rb') as f:
self.mc = Machine(f)
self.foutc_str = f.read().strip().decode()
out4file = os.path.join(curdir, 'lattice/out4_0.lat')
self.out4file = make_latfile(out4file)
with open(self.out4file, 'rb') as f:
self.fout4_str = f.read().strip().decode()
self.latfile4 = os.path.join(curdir, 'lattice/out4_1.lat')
def test_run_2(self):
""" test_run_2: propagate from the first to last, monitor all BPMs
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
obs = fm.get_index_by_type(type='bpm')['bpm']
r0 = m0.propagate(s0, 0, len(m0), observe=obs)
r,s = fm.run(monitor=obs)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_set_machine(self):
fm_none = ModelFlame()
self.assertIsNone(fm_none.machine)
with open(self.testfile, 'rb') as f:
m = Machine(f)
fm_none.machine = m
self.assertEqual(fm_none.machine, m)
def conf_update(machine):
"""Update machine conf() by using current settings.
Returns
-------
FLAME machine object
"""
m = machine
try:
mconf = m.conf()
except:
_LOGGER.error("Failed to load FLAME machine object.")
return None
mc_src = m.conf(m.find(type='source')[0])
for i in range(len(m)):
elem_i = m.conf(i)
ename, etype = elem_i['name'], elem_i['type']
ki = elem_i.keys()
elem_k = set(ki).difference(mc_src.keys())
if etype == 'source':
elem_k.add('vector_variable')
elem_k.add('matrix_variable')
if etype == 'stripper':
elem_k.add('IonChargeStates')
elem_k.add('NCharge')
for k in elem_k:
mconf['elements'][i][k] = m.conf(i)[k]
new_m = Machine(mconf)
return new_m
def test_run_7(self):
""" test_run_7: optional monitor setting 'type'
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
obs = fm.get_index_by_type(type='bpm')['bpm']
r0 = m0.propagate(s0, 0, len(m0), observe=obs)
r,s = fm.run(monitor='bpm')
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_init_with_s1(self):
""" test_init_with_s1: s is not None
"""
with open(self.latfile, 'rb') as f:
m = Machine(f)
s0 = m.allocState({})
s1 = s0.clone()
m.propagate(s1, 0, 1)
ms0 = BeamState(s0)
compare_mstates(self, ms0, s0)
ms1 = BeamState(s0, machine=m)
compare_mstates(self, ms1, s1)
ms1_1 = BeamState(s0, latfile=self.latfile)
compare_mstates(self, ms1_1, s1)
def test_collect_data(self):
""" test_collect_data: get pos, x0, IonEk
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
r0 = m0.propagate(s0, 0, 100, observe=range(100))
data0 = collect_data(r0, pos=True, x0=True, IonEk=True)
fm = ModelFlame(latfile)
r, s = fm.run(from_element=1, to_element=99, monitor=range(100))
data = fm.collect_data(r, pos=True, x0=True, IonEk=True)
self.assertEqual(data0['pos'][1:].tolist(), data['pos'].tolist())
self.assertEqual(data0['x0'][1:].tolist(), data['x0'].tolist())
self.assertEqual(data0['IonEk'][1:].tolist(), data['IonEk'].tolist())
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S=self.M.allocState({})
self.M.propagate(S,0,1)
self.refIonZ = S.ref_IonZ
self.IonZ = S.IonZ # list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = np.zeros((len(self.M),9))
self.LD2 = np.zeros((len(self.M),7))
# store element position data
self.bpmls=[]
self.corls=[]
self.cavls=[]
self.solls=[]
self.cavpara=[]
self.solpara=[]
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm' : self.bpmls.append(i)
elif elem['type'] == 'orbtrim' : self.corls.append(i)
elif elem['type'] == 'rfcavity' :
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'],elem['phi']])
elif elem['type'] == 'solenoid' :
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info','w') as f:
iteminfo=[len(self.M),self.bpmls,self.corls,self.cavls,self.solls]
cPickle.dump(iteminfo,f)
def __init__(self, source, output=None, auto_scaling=False , starting_offset=0.0, **kws):
self._source = source
self._auto_scaling = auto_scaling
self._starting_offset = starting_offset
if type(self._source) == str:
with open(self._source, 'rb') as lat:
self.M = Machine(lat)
elif type(self._source) == Machine:
self.M = self._source
else:
raise ValueError('source must be a file path of .lat or flame.Machine object')
self.types = {'rfcavity': {'flag':True, 'name':'rfcavity', 'color':'orange', 'scale':0.0},
'solenoid': {'flag':True, 'name':'solenoid', 'color':'red', 'scale':0.0},
'quadrupole': {'flag':True, 'name':'quad', 'color':'purple', 'scale':0.0},
'sextupole': {'flag':True, 'name':'sext', 'color':'navy', 'scale':0.0},
'sbend': {'flag':True, 'name':'bend', 'color':'green', 'scale':0.0},
'equad': {'flag':True, 'name':'e-quad', 'color':'blue', 'scale':0.0},
'edipole': {'flag':True, 'name':'e-dipole', 'color':'lime', 'scale':0.0},
'bpm': {'flag':True, 'name':'bpm', 'color':'m', 'scale':0.0},
'orbtrim': {'flag':True, 'name':'corr', 'color':'black', 'scale':0.0},
'stripper': {'flag':True, 'name':'stripper', 'color':'y', 'scale':0.0},
'marker': {'flag':True, 'name':'pm', 'color':'c', 'scale':0.0}
}
if self._auto_scaling:
for i in range(len(self.M)):
elem = self.M.conf(i)
if elem['type'] in self.types.keys():
prv_scl = self.types[elem['type']]['scale']
tmp_scl = np.abs(self._get_scl(elem))
self.types[elem['type']]['scale'] = max(prv_scl,tmp_scl)
if isinstance(output, str):
self.generate(legend=True, option=True)
self.output(window=True, fname=output, **kws)
def test_configure(self):
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
e_cor_idx = 10
m0.reconfigure(10, {'theta_x': 0.005})
m0.propagate(s0, 0, 1)
r0 = m0.propagate(s0, 1, -1, range(len(m0)))
fm = ModelFlame(latfile)
e = fm.get_element(index=10)[0]
e['properties']['theta_x'] = 0.005
fm.configure(e)
r, s = fm.run(monitor=range(len(m0)))
rs0 = [ts for (ti, ts) in r0]
rs = [ts for (ti, ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_reconfigure(self):
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
e_cor_idx = 10
e_name = m0.conf(e_cor_idx)['name']
m0.reconfigure(10, {'theta_x': 0.005})
r0 = m0.propagate(s0, 0, len(m0), range(len(m0)))
fm = ModelFlame(latfile)
fm.reconfigure(e_name, {'theta_x': 0.005})
r, s = fm.run(monitor=range(len(m0)))
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_run_8(self):
""" test_run_8: include_initial_state
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
m0.propagate(s0, 0, 1)
r0 = m0.propagate(s0, 1, len(m0), observe=range(len(m0)))
r,s = fm.run(monitor='all', include_initial_state=False)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def machine_setter(_latfile=None, _machine=None, _handle_name=None):
""" set flame machine, prefer *_latfile*
:return: FLAME machine object
"""
if _latfile is not None:
try:
m = Machine(open(_latfile, 'r'))
except:
if _machine is None:
_LOGGER.error("{}: Failed to initialize flame machine".format(
_handle_name))
return None
else:
_LOGGER.warning("{}: Failed to initialize flame machine, "
"use _machine instead".format(_handle_name))
m = _machine
else:
if _machine is None:
return None
else:
m = _machine
return m
def test_run_4(self):
""" test_run_4: run and monitor from element index of 10 to 20
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
m0.propagate(s0, 0, 10)
r0 = m0.propagate(s0, 10, 11, observe=range(10, 21))
fm = ModelFlame(latfile)
r, s = fm.run(from_element=10, to_element=20, monitor=range(10,21))
compare_mstates(self, s0, s)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_run_5(self):
""" test_run_5: using BeamState object
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
ms = BeamState(machine=m0)
fm = ModelFlame()
fm.bmstate = ms
fm.machine = m0
obs = fm.get_index_by_type(type='bpm')['bpm']
r,s = fm.run(monitor=obs)
s0 = m0.allocState({})
m0.propagate(s0, 0, 1)
r0 = m0.propagate(s0, 1, len(m0), observe=obs)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
class VAF:
"""
VAF
===
Tool box for FLAME python interface
You can start VAF by
>>> va = vaf.VAF(fname=${lattice_file_path})
Initial beam parameters
(default is the same as lattice file)
-------------------------------------
BC0: 7*n array (n: number of charge state)
Barycenter vectors for each charge state.
contains:[x, px, y, py, z, pz, 1]
unit: [mm, rad, mm, rad, rad, MeV/u, 1]
ENV0: 7*7*n array (n: number of charge state)
Envelope matrixes for each charge state.
IonZ: n array (n: number of charge state)
Charge to mass ratio for each charge state.
RefIonEk: float
Reference kinetic energy at starting point [eV/u]
This tool box contains
(Check each documentation for detail.)
--------------------------------------
prt_beam_prms()
prt_lat_list()
getelem(int)
getvalue(int,str)
getindex(str)
getindexu(str)
setvalue(int,str,float)
tcs()
tcs_seg()
SaveBeam(S)
On developping
--------------
genRN()
"""
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S=self.M.allocState({})
self.M.propagate(S,0,1)
self.refIonZ = S.ref_IonZ
self.IonZ = S.IonZ # list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = np.zeros((len(self.M),9))
self.LD2 = np.zeros((len(self.M),7))
# store element position data
self.bpmls=[]
self.corls=[]
self.cavls=[]
self.solls=[]
self.cavpara=[]
self.solpara=[]
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm' : self.bpmls.append(i)
elif elem['type'] == 'orbtrim' : self.corls.append(i)
elif elem['type'] == 'rfcavity' :
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'],elem['phi']])
elif elem['type'] == 'solenoid' :
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info','w') as f:
iteminfo=[len(self.M),self.bpmls,self.corls,self.cavls,self.solls]
cPickle.dump(iteminfo,f)
def getelem(self,num):
"""
getelem(a)
Get base parameter of lattice element.
Parameters
----------
a : int
Input index of the element.
Return
------
out: dict
Python dict style lattice element.
Examples
--------
>>> print va.getelem(8)
OrderedDict([('B', 5.34), ('L', 0.1), ('aper', 0.02), ('name', 'ls1_ca01_sol1_d1131_1'), ('type', 'solenoid')])
"""
#return self.M.conf()['elements'][num]
return self.lat[num]
def getvalue(self,num,name):
"""
getvalue(a, b)
Get latest parameter of lattice element.
Parameters
----------
a: int
Input index of the element
b: str
Input name of the parameter
Return
------
out: depends on parameter name
Value of the parameter.
Example
-------
>>> print va.getvalue(8,'B')
5.34
"""
return self.M.conf(num)[name]
def getindex(self,name,searchfrom='name'):
"""
getindex(a, searchby='name')
Get index list of lattice elements by python style regular expression.
Parameters
----------
a: str
Keyword for search.
searchfrom: str
Search from 'name' or 'type'.
Return
------
out: list
Search result.
"""
name = name.replace(':','_').lower()
pat = re.compile(name)
result = []
for (i,elem) in enumerate(self.lat):
if pat.search(elem[searchby]):
result.append(i)
return result
def getindexu(self,name,searchby='name'):
"""
getindex(a, searchby='name')
Get index list of lattice elements by unix style regular expression.
Parameters
----------
a: str
Keyword for search.
searchfrom: str
Search from 'name' or 'type'.
Return
------
out: list
Search result.
"""
name = name.replace(':','_').lower()
result = []
for (i,elem) in enumerate(self.lat):
if fnmatch.fnmatch(elem[searchby],name):
result.append(i)
return result
def setvalue(self,num,name,val):
"""
setvalue(a, b, c)
Set parameter of lattice element.
Parameters
----------
a: int
Input index of the element.
b: str
Input name of the parameter.
c: float
Set value to the parameter.
"""
self.M.reconfigure(num,{name:float(val)})
# Developing
def genRN(self,scl=0.0005,seed=None):
np.random.seed(seed)
for (i,elem) in enumerate(self.lat):
rnds = np.random.randn(5)*scl #rms 0.5 mm and rms 0.5 mrad
etype = elem['type']
if etype == 'rfcavity' or etype == 'solenoid' or etype == 'quadrupole':
self.M.reconfigure(i,{'dx':float(rnds[0])})
self.M.reconfigure(i,{'dy':float(rnds[1])})
#self.M.reconfigure(i,{'pitch':float(rnds[2])})
#self.M.reconfigure(i,{'yaw':float(rnds[3])})
#self.M.reconfigure(i,{'tilt':float(rnds[4])})
def prt_beam_prms(self):
"""
Print initial beam parameters in lattice file
"""
print '\nIon Charge States = ', self.M.conf()['IonChargeStates']
print 'IonEs [MeV] = ', self.M.conf()['IonEs']/1e6
print 'IonEk [MeV] = ', self.M.conf()['IonEk']/1e6
print '\nBaryCenter 0:\n', self.M.conf()['BaryCenter0']
print '\nBaryCenter 1:\n', self.M.conf()['BaryCenter1']
print '\nBeam Envelope 0:\n', self.M.conf()['S0']
print '\nBeam Envelope 1:\n', self.M.conf()['S1']
def prt_lat_list(self):
"""
Print all lattice elements with index.
returns (index, name, type)
"""
for (i,elem) in enumerate(self.lat):
print(i, elem['name'], elem['type'])
def tcs(self,lpf=0, opf=1):
"""
tcs(lpf=0, opf=1)
Main function for one through calculation.
Calculation data is stored as VAF.LD.
VAF.LD[i]:
----------
i: int
index of the element
contains:
[pos, xcen, ycen, zcen, xrms, yrms, zrms, ref_phis, ref_IonEk]
units:
[m, mm, mm, mm, mm, mm, mm, rad, eV/u]
Parameters
----------
lpf: bool
flag for loop run.
opf: bool
flag for file type output.
Returns
-------
S: object
Beam parameters at the ending point.
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
#S.clng = self.clng
fin = len(self.M)
# store initial beam data
self.LD[0][0] = S.pos
#Mean data
self.LD[0][1] = S.moment0_env[0]
self.LD[0][2] = S.moment0_env[2]
self.LD[0][3] = S.moment0_env[4]
self.LD[0][4] = S.moment0_rms[0]
self.LD[0][5] = S.moment0_rms[2]
self.LD[0][6] = S.moment0_rms[4]
self.LD[0][7] = S.ref_phis
self.LD[0][8] = S.ref_IonEk
# store initial beam data
self.LD2[0][0] = S.pos
#Mean data
self.LD2[0][1] = S.moment0_env[1]
self.LD2[0][2] = S.moment0_env[3]
self.LD2[0][3] = S.moment0_env[5]
self.LD2[0][4] = S.moment0_rms[1]
self.LD2[0][5] = S.moment0_rms[3]
self.LD2[0][6] = S.moment0_rms[5]
# propagate step by step and store beam data
for i in range(1,len(self.M)):
self.M.propagate(S, i, 1)
self.LD[i][0] = S.pos
#Mean data
self.LD[i][1] = S.moment0_env[0]
self.LD[i][2] = S.moment0_env[2]
self.LD[i][3] = S.moment0_env[4]
self.LD[i][4] = S.moment0_rms[0]
self.LD[i][5] = S.moment0_rms[2]
self.LD[i][6] = S.moment0_rms[4]
self.LD[i][7] = S.ref_phis
self.LD[i][8] = S.ref_IonEk
self.LD2[i][0] = S.pos
#Mean data
self.LD2[i][1] = S.moment0_env[1]
self.LD2[i][2] = S.moment0_env[3]
self.LD2[i][3] = S.moment0_env[5]
self.LD2[i][4] = S.moment0_rms[1]
self.LD2[i][5] = S.moment0_rms[3]
self.LD2[i][6] = S.moment0_rms[5]
#output data for plotting
if opf: np.savetxt('ldata.txt',self.LD)
if not lpf: return S
def tcs2(self):
"""
tcs(lpf=0, opf=1)
Main function for one through calculation.
Calculation data is stored as VAF.LD.
VAF.LD[i]:
----------
i: int
index of the element
contains:
[pos, xcen, ycen, zcen, xrms, yrms, zrms, ref_phis, ref_IonEk]
units:
[m, mm, mm, mm, mm, mm, mm, rad, eV/u]
Parameters
----------
lpf: bool
flag for loop run.
opf: bool
flag for file type output.
Returns
-------
S: object
Beam parameters at the ending point.
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
### Reconstract data
U = collections.OrderedDict()
T = collections.OrderedDict()
U,T = self.save(U, T, S)
#S.clng = self.clng
fin = len(self.M)
H = self.M.propagate(S, 1, fin, observe=range(fin))
### Reconstract data
U = collections.OrderedDict()
T = collections.OrderedDict()
for i in range(fin-1):
U,T = self.save(U, T, H[i][1])
return U,T
def tcss(self,start=None,end=None,S_in=None,obs=None):
if start == None : start = 1
if end == None : end = len(self.M) -1
if S_in == None :
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
else :
# set initial beam data
S = S_in.clone()
if obs == None :
obs = range(len(self.M))
#S.clng = self.clng
H_ini = (0, S.clone())
fin = end - start + 1
#H = self.M.propagate(S, start, fin, observe=range(len(self.M)))
H = self.M.propagate(S, start, fin, observe=obs)
if not obs[0]: return [H_ini] + H
else: return H
def looptcs(self):
"""
Main loop for Virtual Accelerator
This function should be called as background job.
Loop is stopped by setting itr = 1.
"""
while self.itr < 1:
#self.genRandomNoise() #developing
self.tcs(lpf=1)
#self.itr +=1
def save(self,u, t, S, label = 'total'):
for (k,j) in enumerate(S.IonZ):
cs = str(int(round(j*238.0)))
try:
u[cs].pos.append(S.pos)
except:
u[cs] = self.SPI()
u[cs].pos.append(S.pos)
for i,(n,m) in enumerate(zip(cen,rms)):
getattr(u[cs], n).append(S.moment0[i,k])
getattr(u[cs], m).append(np.sqrt(S.moment1[i,i,k]))
try:
t[label].pos.append(S.pos)
except:
t[label] = self.SPI()
t[label].pos.append(S.pos)
for i,(n,m) in enumerate(zip(cen,rms)):
getattr(t[label], n).append(S.moment0_env[i])
getattr(t[label], m).append(S.moment0_rms[i])
t[label].ek.append(S.ref_IonEk)
t[label].phis.append(S.ref_phis)
return u,t
class SaveBeam:
"""
SaveBeam(S_in)
Store beam parameters for VAF.tcs_seg().
Parameters
----------
S_in: object
Beam parameters made by VAF.tcs() or VAF.tcs_seg().
Returns
-------
out: object
Beam parameters object for VAF.tcs_seg().
"""
def __init__(self,S_in):
self.refIonZ = S_in.ref_IonZ
self.IonZ_io = S_in.IonZ
self.refphis = S_in.ref_phis
self.phis_io = S_in.phis
self.BC0 = S_in.moment0
self.ENV0 = S_in.moment1
self.refIonEk = S_in.ref_IonEk
self.pos = S_in.pos
class SPI:
def __init__(self):
self.pos = []
self.x = []
self.xp = []
self.y = []
self.yp = []
self.z = []
self.zp = []
self.xrms = []
self.xprms = []
self.yrms = []
self.yprms = []
self.zrms = []
self.zprms = []
self.ek = []
self.phis = []
def n(self, name):
self.x_n = np.array(self.x)
#---------------------------------------------------------
def rms2(self,xy,S,cs=0):
if xy == 'x': i=0
elif xy == 'px': i=1
elif xy == 'y': i=2
elif xy == 'py': i=3
elif xy == 'z': i=4
elif xy == 'pz': i=5
return S.moment1[i,i,cs]
def rms(self,xy,S,cs=0):
return np.sqrt(self.rms2(xy,S,cs))
def emit(self,xy,S,cs=0):
if xy == 'x': i=0
elif xy == 'y': i=2
else : raise NameError('First argument must be \'x\' or \'y\'')
xx = S.moment1[i,i,cs]*1e-6
xxp = S.moment1[i,i+1,cs]*1e-3
xpxp = S.moment1[i+1,i+1,cs]
return np.sqrt(np.abs(xx*xpxp - xxp*xxp))
def nemit(self,xy,S,cs=0):
bg = S.beta[cs]*S.gamma[cs]
return bg*self.emit(xy,S,cs)
def bfunc(self,xy,S,cs=0):
return self.rms2(xy,S,cs)*1e-6/self.emit(xy,S,cs)
def bfuncp(self,xy,S,cs=0):
if xy == 'x': i = 0
elif xy == 'y': i = 2
return S.moment1[i,i+1,cs]*1e-3/self.emit(xy,S,cs)
def hbfunc(self,xy,f,cs=0):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = self.bfunc(xy,f[i][1],cs)
return ans
def hbfuncp(self,xy,f,cs=0):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = self.bfuncp(xy,f[i][1],cs)
return ans
def hpos(self,f):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = f[i][1].pos
return ans
def hcen(self,xy,f,cs=0):
if xy == 'x': i=0
elif xy == 'px': i=1
elif xy == 'y': i=2
elif xy == 'py': i=3
elif xy == 'z': i=4
elif xy == 'pz': i=5
ans = np.zeros(len(f))
if cs == -1:
for k in range(len(f)):
ans[k] = f[k][1].moment0_env[i]
else:
for k in range(len(f)):
ans[k] = f[k][1].moment0[i,cs]
return ans
def hrms(self,xy,f,cs=0):
if xy == 'x': i=0
elif xy == 'px': i=1
elif xy == 'y': i=2
elif xy == 'py': i=3
elif xy == 'z': i=4
elif xy == 'pz': i=5
ans = np.zeros(len(f))
if cs == -1:
for k in range(len(f)):
ans[k] = f[k][1].moment0_rms[i]
else:
for k in range(len(f)):
ans[k] = self.rms(xy,f[k][1],cs)
return ans
def dpp(self,S,cs=0):
gm = S.gamma[cs]
iw = S.IonEk[cs]
dw = S.moment0[5][cs]*1e6
return gm/(1+gm)*dw/iw
#---------------------------------------------------------
def hdpp(self,f,cs=0):
ans = np.zeros(len(f))
for k in range(len(f)):
ans[k] = self.dpp(f[k][1],cs=cs)
return ans
def hdsp(self,xy,f,cs=0):
if xy == 'x' : i = 0
elif xy == 'y' : i = 2
ans = np.zeros(len(f))
for k in range(len(f)):
gm = f[k][1].gamma[cs]
iw = f[k][1].IonEk[cs]
xw = f[k][1].moment1[i,5,cs]*1e-3
ww = f[k][1].moment1[5,5,cs]*1e6*gm/(1+gm)/iw
ans[k] = xw/ww
return ans
def getek(self,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].ref_IonEk
return ans
def getid0(self,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].moment0_env[id]
return ans
def getid00(self,cs,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].moment0[id,cs]
return ans
def getid1(self,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].moment0_rms[id]
return ans
def getid01(self,cs,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = np.sqrt(f[j][1].moment1[id,id,cs])
return ans
#!/usr/bin/env python
import sys
"[includes]"
from flame import Machine
"[includes]"
"[Parse lattice file]"
with open(sys.argv[1], "r") as F:
mymachine = Machine(F)
"[Parse lattice file]"
"[Allocate State]"
thestate = mymachine.allocState({})
"[Allocate State]"
print "Initial state", thestate
"[Run simulation]"
mymachine.propagate(thestate)
"[Run simulation]"
print "Final state", thestate
class VAF:
"""
Contains:
getelem(int)
getvalue(int,str)
getindex(str)
getindexu(str)
setelem(int,str,float)
prt_lat_prms()
prt_lat()
tcs()
tcs(int,int)
looptcs()
SaveBeam(S)
"""
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S=self.M.allocState({})
self.M.propagate(S,0,1)
self.refIonZ = S.ref_IonZ
self.IonZ_io = S.IonZ# list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = numpy.zeros((len(self.M),9))
# store element position data
self.bpmls=[]
self.corls=[]
self.cavls=[]
self.solls=[]
self.cavpara=[]
self.solpara=[]
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm' : self.bpmls.append(i)
elif elem['type'] == 'orbtrim' : self.corls.append(i)
elif elem['type'] == 'rfcavity' :
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'],elem['phi']])
elif elem['type'] == 'solenoid' :
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info','w') as f:
iteminfo=[len(self.M),self.bpmls,self.corls,self.cavls,self.solls]
cPickle.dump(iteminfo,f)
def getelem(self,num):
"""
Get parameter of lattice element
getelem(index of lattice element)
"""
print self.lat[num]
def getvalue(self,num,name):
"""
Get parameter of lattice element
getelem(index of lattice element, parameter name)
"""
print self.M.conf(num)[name]
def getindex(self,name,searchby='name'):
"""
Get index list of lattice elements by python style regular expression
getindex(name or type,
searchby = 'name')
"""
name = name.replace(':','_').lower()
pat = re.compile(name)
result = []
for (i,elem) in enumerate(self.lat):
if pat.search(elem[searchby]):
result.append(i)
return result
def getindexu(self,name,searchby='name'):
"""
Get index list of lattice elements by unix style regular expression
getindex(name or type,
searchby = 'name')
"""
name = name.replace(':','_').lower()
result = []
for (i,elem) in enumerate(self.lat):
if fnmatch.fnmatch(elem[searchby],name):
result.append(i)
return result
def setelem(self,num,name,val):
"""
Set parameter of lattice element
setelem(position number of lattice element,
name of parameter
value of parameter)
"""
#D = self.M.conf()['elements'][num]
#D = self.lat[num]
#D[name] = float(val)
self.M.reconfigure(num,{name:float(val)})
"""
# Developing
def genRandomNoise(self):
for (i,cv) in enumerate(self.cavls):
D = self.M.conf()['elements'][cv]
for (j,st) in enumerate(['scl_fac','phi']):
defv = self.cavpara[i][j]
D[st] = defv + numpy.random.normal(0.0,abs(defv)*1e-3)
self.M.reconfigure(cv, D)
D.clear()
for (i,cv) in enumerate(self.solls):
D = self.M.conf()['elements'][cv]
for (j,st) in enumerate(['B']):
defv = self.solpara[i][j]
D[st] = defv + numpy.random.normal(0.0,abs(defv)*1e-3)
self.M.reconfigure(cv, D)
D.clear()
def loopRN(self):
while self.itr < 100:
self.genRandomNoise()
self.itr += 1
"""
def prt_lat_prms(self):
"""Print initial beam parameters in lattice file"""
print '\nIon Charge States = ', self.M.conf()['IonChargeStates']
print 'IonEs [MeV] = ', self.M.conf()['IonEs']/1e6
print 'IonEk [MeV] = ', self.M.conf()['IonEk']/1e6
print '\nBaryCenter 0:\n', self.M.conf()['BaryCenter0']
print '\nBaryCenter 1:\n', self.M.conf()['BaryCenter1']
print '\nBeam Envelope 0:\n', self.M.conf()['S0']
print '\nBeam Envelope 1:\n', self.M.conf()['S1']
def prt_lat(self):
"""Print all lattice elements"""
for (i,elem) in enumerate(self.lat):
print(i, elem['name'], elem['type'])
def tcs(self,lpf=0, opf=1):
"""
Main function of one through calculation
Calculation data is stored at LD.
LD[i] contains:
[pos, xcen, ycen, zrad, xrms, yrms, ref_phis, ref_IonEk]
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ_io
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
#S.clng = self.clng
fin = len(self.M)
# store initial beam data
self.LD[0][0] = S.pos
#Mean data
self.LD[0][1] = S.moment0_env[0]
self.LD[0][2] = S.moment0_env[2]
self.LD[0][3] = S.moment0_env[4]
self.LD[0][4] = S.moment0_rms[0]
self.LD[0][5] = S.moment0_rms[2]
self.LD[0][6] = S.moment0_rms[4]
self.LD[0][7] = S.ref_phis
self.LD[0][8] = S.ref_IonEk
# propagate step by step and store beam data
for i in range(1,len(self.M)):
self.M.propagate(S, i, 1)
self.LD[i][0] = S.pos
#Mean data
self.LD[i][1] = S.moment0_env[0]
self.LD[i][2] = S.moment0_env[2]
self.LD[i][3] = S.moment0_env[4]
self.LD[i][4] = S.moment0_rms[0]
self.LD[i][5] = S.moment0_rms[2]
self.LD[i][6] = S.moment0_rms[4]
self.LD[i][7] = S.ref_phis
self.LD[i][8] = S.ref_IonEk
#output data for plotting
if opf: numpy.savetxt('ldata.txt',self.LD)
if not lpf: return S
def tcs_seg(self,start,end,SD=None,allcs=0,opf=1):
"""
Main function of segmented section calculation
(S, RD) = tcs_seg(start_index, end_index, SD=beam_data)
beam_data can be generated by SaveBeam(S)
RD[i] contains:
[pos, xcen, ycen, zrad, xrms, yrms, ref_phis, ref_IonEk]
i: index of the element
For all charge outputs: allcs=1,
(S, RD, AD) = tcs_seg(start_index, end_index, SD=beam_data, allcs=1)
AD[k][i] contains:
[pos, xcen, ycen, zrad, xrms, yrms]
k: index of the charge state
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
if SD == None :
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ_io
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
else :
# set initial beam data
S.ref_IonZ = SD.refIonZ
S.IonZ = SD.IonZ_io
S.moment0 = SD.BC0
S.moment1 = SD.ENV0
S.ref_phis = SD.refphis
S.phis = SD.phis_io
S.ref_IonEk = SD.refIonEk
S.IonEk = SD.BC0[PS_PS,:]*MeVtoeV + SD.refIonEk
S.pos = SD.pos
phis_ini = S.ref_phis
#S.clng = self.clng
fin = end - start + 1
RD = numpy.zeros((fin,8))
#if allcs: AD = [[[0.0]*6 for i in range(fin)] for j in range(len(S.IonZ))]
if allcs: AD = numpy.zeros((len(S.IonZ),fin,8))
# store initial beam data
RD[0][0] = S.pos
RD[0][1] = S.moment0_env[0]
RD[0][2] = S.moment0_env[2]
RD[0][3] = S.moment0_env[4]
RD[0][4] = S.moment0_rms[0]
RD[0][5] = S.moment0_rms[2]
RD[0][6] = S.ref_phis - phis_ini
RD[0][7] = S.ref_IonEk
if allcs:
for k in range(len(S.IonZ)):
AD[k][0][0] = S.pos
AD[k][0][1] = S.moment0[0,k]
AD[k][0][2] = S.moment0[2,k]
AD[k][0][3] = S.moment0[4,k]
AD[k][0][4] = numpy.sqrt(S.moment1[0,0,k])
AD[k][0][5] = numpy.sqrt(S.moment1[2,2,k])
# propagate step by step and store beam data
for (j,i) in enumerate(range(start,end)):
self.M.propagate(S, i+1, 1)
RD[j+1][0] = S.pos
RD[j+1][1] = S.moment0_env[0]
RD[j+1][2] = S.moment0_env[2]
RD[j+1][3] = S.moment0_env[4]
RD[j+1][4] = S.moment0_rms[0]
RD[j+1][5] = S.moment0_rms[2]
RD[j+1][6] = S.ref_phis - phis_ini
RD[j+1][7] = S.ref_IonEk
if allcs:
for k in range(len(S.IonZ)):
AD[k][j+1][0] = S.pos
AD[k][j+1][1] = S.moment0[0,k]
AD[k][j+1][2] = S.moment0[2,k]
AD[k][j+1][3] = S.moment0[4,k]
AD[k][j+1][4] = numpy.sqrt(S.moment1[0,0,k])
AD[k][j+1][5] = numpy.sqrt(S.moment1[2,2,k])
if opf: numpy.savetxt('ldata.txt',RD)
if allcs: return (S,RD,AD)
else : return (S,RD)
def looptcs(self):
"""
Main loop for Virtual Accelerator
This function should be called as background job.
Loop is stopped by setting itr = 1.
"""
while self.itr < 1:
#self.genRandomNoise() #developing
self.tcs(lpf=1)
#self.itr +=1
class SaveBeam:
def __init__(self,S_in):
self.refIonZ = S_in.ref_IonZ
self.IonZ_io = S_in.IonZ
self.refphis = S_in.ref_phis
self.phis_io = S_in.phis
self.BC0 = S_in.moment0
self.ENV0 = S_in.moment1
self.refIonEk = S_in.ref_IonEk
self.pos = S_in.pos
