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_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})
r0 = m0.propagate(s0, 0, len(m0), 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)
class TestGenerateLatfile(unittest.TestCase):
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 tearDown(self):
for f in [self.latfile0, self.latfile1, self.latfile2, self.latfile4]:
if os.path.isfile(f):
os.remove(f)
def test_generate_latfile_original1(self):
sio = StringIO()
sioname = generate_latfile(self.m, out=sio)
self.assertEqual(sioname, 'string')
self.assertEqual(sio.getvalue().strip(), self.fout1_str)
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())
@unittest.skip("output precision depends on the python version")
def test_generate_latfile_original3(self):
sio = StringIO()
sioname = generate_latfile(self.mc, original=self.testcfile, out=sio)
self.assertEqual(sioname, 'string')
self.assertEqual(sio.getvalue().strip(), self.fout3_str)
def test_generate_latfile_original4(self):
sio = StringIO()
sioname = generate_latfile(self.m, start=30, end=60, out=sio)
self.assertEqual(sioname, 'string')
self.assertEqual(sio.getvalue().strip(), self.fout4_str)
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])
},
'LS1_CB03:DCV_D1403': {
'config': {
'theta_y': 0.00028278723377
},
'id': 250
},
'LS1_CB03:DCV_D1363': {
'config': {
'theta_y': 0.00035444591542
},
'id': 220
},
'LS1_CB06:DCV_D1574': {
'config': {
'theta_y': 0.00013932249997
},
'id': 393
},
'LS1_CB01:DCV_D1255': {
'config': {
'theta_y': 0.0003313366176
},
'id': 131
}
}
for k, v in opt_cor.items():
m.reconfigure(v['id'], v['config'])
flameutils.generate_latfile(m, 'optout.lat')
def _co_execute(self):
"""Execute the virtual accelerator. This includes the following:
1. Creating a temporary working directory for execution of FLAME.
2. Set up the working directory by symlinking from the data directory.
3. Writing the EPICS DB to the working directory (va.db).
4. Starting the softIoc and channel initializing monitors.
5. Add noise to the settings for all input (CSET) channels.
6. Create or update the FLAME machine configuration.
7. Propagate the FLAME simulation and read the results.
8. Update the READ channels of all devives.
9. Update the REST channels of input devies.
10. Repeat from step #5.
"""
_LOGGER.debug("VA: Execute virtual accelerator")
_LOGGER.info("VA: Running at " + self._ts_now)
chanprefix = self._chanprefix
if self._pv_suffix != '':
suffix_str = "_" + self._pv_suffix
else:
suffix_str = ""
# Add channel for VA rep-rate
chanrate = f"{chanprefix}:SVR:RATE{suffix_str}"
self._epicsdb.append(("ao", chanrate, OrderedDict([
("DESC", "Rep-rate of Simulation Engine"),
("VAL", self._rate),
("PREC", 1),
])))
_LOGGER.info("VA: Reprate PV is " + chanrate)
# Add channel for sample counting
chansample_cnt = f"{chanprefix}:SVR:CNT{suffix_str}"
self._epicsdb.append(("ao", chansample_cnt, OrderedDict([
("DESC", "Sample counter for scan client"),
("VAL", 0)
])))
# Add channel for VA configuration and control
channoise = f"{chanprefix}:SVR:NOISE{suffix_str}"
self._epicsdb.append(("ao", channoise, OrderedDict([
("DESC", "Noise level of Virtual Accelerator"),
("VAL", self._noise),
("PREC", 5)
])))
_LOGGER.info("VA: Noise PV is " + channoise)
chanstat = f"{chanprefix}:SVR:STATUS{suffix_str}"
self._epicsdb.append(("bi", chanstat, OrderedDict([
("DESC", "Status of Virtual Accelerator"),
("VAL", 1),
("ZNAM", "ERR"),
("ONAM", "OK"),
("PINI", "1")
])))
_LOGGER.info("VA: Status PV is " + chanstat)
# MPS status
chan_mps_stat = f"{chanprefix}:SVR:MpsStatus"
self._epicsdb.append(("mbbi", chan_mps_stat, OrderedDict([
("DESC", "MPS Status of Virtual Accelerator"),
("VAL", 3),
("ZRST", "Fault"),
("ONST", "Disable"),
("TWST", "Monitor"),
("THST", "Enable"),
("PINI", "1")
])))
_LOGGER.info("VA: MPS PV is " + chan_mps_stat)
#
chancharge = f"{chanprefix}:SVR:CHARGE{suffix_str}"
self._epicsdb.append(("ai", chancharge, OrderedDict([
("DESC", "Q/M of Virtual Accelerator"),
("VAL", 0.0),
("PREC", 5)
])))
# initial beam condition
chanbsrc = f"{chanprefix}:SVR:BEAM{suffix_str}"
self._epicsdb.append(("waveform", chanbsrc, OrderedDict([
("DESC", "Init beam of Virtual Accelerator"),
("NELM", 4096),
("FTVL", "UCHAR")
])))
_LOGGER.info("VA: Init beam condition PV is " + chanbsrc)
#
if self.work_dir is not None:
os.makedirs(self.work_dir)
self._rm_work_dir = False
latticepath = os.path.join(self.work_dir, "test.lat")
else:
self.work_dir = tempfile.mkdtemp(_TEMP_DIRECTORY_SUFFIX)
self._rm_work_dir = True
latticepath = None
_LOGGER.info("VA: Working directory: %s", self._work_dir)
# input file paths
epicsdbpath = os.path.join(self.work_dir, "va.db")
#output file paths
epicslogpath = os.path.join(self.work_dir, "softioc.log")
if os.path.isabs(self.data_dir):
abs_data_dir = self.data_dir
self._latfactory.dataDir = self.data_dir
else:
abs_data_dir = os.path.abspath(self.data_dir)
self._latfactory.dataDir = os.path.abspath(self.data_dir)
with open(epicsdbpath, "w") as outfile:
self._write_epicsdb(outfile)
_LOGGER.info("VA: Write EPICS database to %s", epicsdbpath)
#_LOGGER.debug("VA: Write EPICS database to %s", epicsdbpath)
self._ioc_logfile = open(epicslogpath, "w")
self._ioc_process = Popen(["softIoc", "-d", "va.db"], cwd=self.work_dir,
stdout=self._ioc_logfile, stderr=subprocess.STDOUT)
_LOGGER.debug("VA: Start EPICS soft IOC with log %s", epicslogpath)
_LOGGER.debug("VA: Connecting to channels: {}".format(len(self._csetmap.keys())))
self._subscriptions = []
self._subscriptions.append(catools.camonitor(chanrate, self._handle_rate_monitor))
self._subscriptions.append(catools.camonitor(channoise, self._handle_noise_monitor))
self._subscriptions.append(catools.camonitor(chanbsrc, self._handle_bsrc_monitor))
self._subscriptions.extend(catools.camonitor(self._csetmap.keys(), self._handle_cset_monitor))
_LOGGER.debug("VA: Connecting to channels: Done")
machine = None
while self._continue:
# update the RSET channels with new settings
batch = catools.CABatch()
for cset in self._csetmap.items():
name, field = self._fieldmap[cset[0]]
batch[cset[1][0]] = self._settings[name][field]
batch.caput()
settings = self._copy_settings_with_noise()
self._latfactory.settings = settings
lattice = self._latfactory.build()
start = time.time()
if machine is None:
_LOGGER.debug("VA: Create FLAME machine from configuration")
machine = Machine(lattice.conf())
else:
_LOGGER.debug("VA: Reconfigure FLAME machine from configuration")
for idx, elem in enumerate(lattice.elements):
machine.reconfigure(idx, elem[2])
if self._bsrc is not None:
_LOGGER.info("VA: Reconfigure FLAME machine with init beam config")
machine.reconfigure(self._bsrc['index'], self._bsrc['properties'])
if latticepath is not None:
_LOGGER.debug(f"VA: Write FLAME lattice file to {outfile}")
generate_latfile(machine, latfile=latticepath)
_LOGGER.debug("VA: Allocate FLAME state from configuration")
S = machine.allocState({})
output_map = []
for elem in lattice.elements:
if 'name' in elem[3]:
output_map.append(elem[3]['name'])
else:
output_map.append(None)
batch = catools.CABatch()
for i in range(0, len(machine)):
machine.propagate(S, i, 1)
if output_map[i] in self._elemmap:
elem = self._elemmap[output_map[i]]
if isinstance(elem, BPMElement):
x_centroid = S.moment0_env[0]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.x_phy], x_centroid)
batch[self._readfieldmap[elem.name][elem.fields.x_phy]] = x_centroid
y_centroid = S.moment0_env[2]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.y_phy], y_centroid)
batch[self._readfieldmap[elem.name][elem.fields.y_phy]] = y_centroid
# convert rad to deg and adjust for 161MHz sampling frequency
phase = _normalize_phase(2.0 * S.ref_phis * (180.0 / math.pi))
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.phase_phy], phase)
batch[self._readfieldmap[elem.name][elem.fields.phase_phy]] = phase
energy = S.ref_IonEk/1.0e6 # convert eV to MeV
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.energy_phy], energy)
batch[self._readfieldmap[elem.name][elem.fields.energy_phy]] = energy
elif isinstance(elem, PMElement):
x_centroid = S.moment0_env[0]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.x], x_centroid)
batch[self._readfieldmap[elem.name][elem.fields.x]] = x_centroid
y_centroid = S.moment0_env[2]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.y], y_centroid)
batch[self._readfieldmap[elem.name][elem.fields.y]] = y_centroid
x_rms = S.moment0_rms[0]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.xrms], x_rms)
batch[self._readfieldmap[elem.name][elem.fields.xrms]] = x_rms
y_rms = S.moment0_rms[2]/1.0e3
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.yrms], y_rms)
batch[self._readfieldmap[elem.name][elem.fields.yrms]] = y_rms
sign = elem.sign
xy_centroid = (sign*x_centroid + y_centroid)/math.sqrt(2.0) # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.xy], xy_centroid)
batch[self._readfieldmap[elem.name][elem.fields.xy]] = xy_centroid
xy_rms = 1.0e-3*math.sqrt(
(S.moment1_env[0, 0] + S.moment1_env[2, 2])*0.5
+ sign*S.moment1_env[0, 2]
)
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.xyrms], xy_rms)
batch[self._readfieldmap[elem.name][elem.fields.xyrms]] = xy_rms
cxy = sign * S.moment1_env[0, 2] * 1e-6 / x_rms / y_rms
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.cxy], cxy)
batch[self._readfieldmap[elem.name][elem.fields.cxy]] = cxy
elif isinstance(elem, (FCElement, VDElement, TargetElement, DumpElement, WedgeElement)):
x_centroid = S.moment0_env[0]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.x], x_centroid)
batch[self._readfieldmap[elem.name][elem.fields.x]] = x_centroid
y_centroid = S.moment0_env[2]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.y], y_centroid)
batch[self._readfieldmap[elem.name][elem.fields.y]] = y_centroid
x_rms = S.moment0_rms[0]/1.0e3 # convert mm to m
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.xrms], x_rms)
batch[self._readfieldmap[elem.name][elem.fields.xrms]] = x_rms
y_rms = S.moment0_rms[2]/1.0e3
_LOGGER.debug("VA: Update read: %s to %s",
self._readfieldmap[elem.name][elem.fields.yrms], y_rms)
batch[self._readfieldmap[elem.name][elem.fields.yrms]] = y_rms
batch.caput()
_LOGGER.info("VA: FLAME execution time: %f s", time.time()-start)
# Allow the BPM, PM, etc. readbacks to update
# before the device setting readbacks PVs.
cothread.Yield()
batch = catools.CABatch()
for name, value in self._csetmap.items():
name, field = self._fieldmap[name]
_LOGGER.debug("VA: Update read: %s to %s", value[1], settings[name][field])
batch[value[1]] = settings[name][field]
batch.caput()
# Sleep for a fraction (10%) of the total execution time
# when one simulation costs more than 0.50 seconds.
# Otherwise, sleep for the rest of 1 second.
# If a scan is being done on this virtual accelerator,
# then the scan server has a period of time to update
# setpoints before the next run of IMPACT.
delt = time.time() - start
if delt > 0.50:
cothread.Sleep(delt * 0.1)
else:
cothread.Sleep(1.0 / self._rate - delt)
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
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
def configure(machine=None, econf=None, **kws):
"""Configure FLAME machine.
Parameters
----------
machine :
FLAME machine object.
econf : (list of) dict
Element configuration(s).
Keyword Arguments
-----------------
latfile : str
FLAME lattice file.
c_idx : int
Element index.
c_dict : dict
Configuration dict.
Returns
-------
m : New FLAME machine object
None if failed, else new machine object.
Note
----
If wanna configure FLAME machine by conventional way, then *c_idx* and
*c_dict* could be used, e.g. reconfigure one corrector of ``m``:
``configure(machine=m, c_idx=10, c_dict={'theta_x': 0.001})``
which is just the same as: ``m.reconfigure(10, {'theta_x': 0.001})``.
Examples
--------
>>> from flame import Machine
>>> from phantasy import flameutils
>>>
>>> latfile = 'test.lat'
>>> m = Machine(open(latfile, 'r'))
>>>
>>> # reconfigure one element
>>> e1 = flameutils.get_element(_machine=m, index=1)[0]
>>> print(e1)
{'index': 1, 'properties': {'L': 0.072, 'aper': 0.02,
'name': 'LS1_CA01:GV_D1124', 'type': 'drift'}}
>>> e1['properties']['aper'] = 0.04
>>> m = flameutils.configure(m, e1)
>>> print(flameutils.get_element(_machine=m, index=1)[0])
{'index': 1, 'properties': {'L': 0.072, 'aper': 0.04,
'name': 'LS1_CA01:GV_D1124', 'type': 'drift'}}
>>>
>>> # reconfiugre more than one element
>>> e_cor = flameutils.get_element(_machine=m, type='orbtrim')
>>> # set all horizontal correctors with theta_x = 0.001
>>> for e in e_cor:
>>> if 'theta_x' in e['properties']:
>>> e['properties']['theta_x'] = 0.001
>>> m = flameutils.configure(m, e_cor)
See Also
--------
get_element : Get FLAME lattice element configuration.
"""
_latfile = kws.get('latfile', None)
_machine = machine
m = machine_setter(_latfile, _machine, 'configure')
if m is None:
return None
_m = Machine(m.conf())
_c_idx, _c_dict = kws.get('c_idx'), kws.get('c_dict')
if _c_idx is not None and _c_dict is not None:
_m.reconfigure(_c_idx, _c_dict)
else:
if not isinstance(econf, list):
_m.reconfigure(econf['index'], econf['properties'])
else:
for e in econf:
_m.reconfigure(e['index'], e['properties'])
return _m
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
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
ax1 = fig1.add_subplot(111)
linex, = ax1.plot(pos[observe_ids],
x[observe_ids],
'r-',
alpha=0.6,
label='$\mathrm{ref\;orbit}$')
linex.set_lw(2)
## apply random kicks
N = 1
#corh_ids_enabled = np.random.choice(corh_ids, size=N)
#corh_val_enabled = 5e-3 * (np.random.random(size=N) * (2 - 1) + 1)
corh_ids_enabled = np.array([392])
corh_val_enabled = np.array([0.005])
for i, v in zip(corh_ids_enabled, corh_val_enabled):
m.reconfigure(i, {'theta_x': v})
"""
for i, v in zip(corh_sel, corh_val):
m.reconfigure(i, {'theta_x': v})
for i, v in zip(corv_sel, corv_val):
m.reconfigure(i, {'theta_y': v})
"""
s_tmp = m.allocState({})
r_tmp = m.propagate(s_tmp, 0, len(m), observe=range(len(m)))
x_tmp, y_tmp = np.array([[r_tmp[i][1].moment0_env[0] for i in range(len(m))]
for j in [0, 2]])
pos = np.array([r_tmp[i][1].pos for i in range(len(m))])
# data plot
linex_tmp, = ax1.plot(pos[observe_ids],
#mymachine = Machine(open('./LS1FS1_latticeE.lat'))
mymachine = Machine(open('./test.lat'))
"[Parse lattice file]"
"[Allocate State]"
thestate = mymachine.allocState({})
"[Allocate State]"
print "Initial state", thestate
"[Run simulation]"
bpms = mymachine.find(type='bpm')
#sel_cor = mymachine.find(name='LS1_CA01:DCH_D1131')[-1]
#mymachine.reconfigure(sel_cor, {'theta_x': 10})
mymachine.reconfigure(9, {'theta_x': 1.00017})
mymachine.reconfigure(10, {'theta_y': 2.00035})
r = mymachine.propagate(thestate, observe=bpms)
"[Run simulation]"
print "Final state", thestate
r01 = r[0][1]
print r01.pos
print r01.moment0_env[0]
print r01.moment0_env[2]
import numpy as np
import matplotlib.pyplot as plt
