Commit bd1af93e authored by Paul Jung's avatar Paul Jung
Browse files

Initial commit

parents
Pipeline #5686 failed with stages
import numpy as np
import math
# rotation xz plane
def rotation_xz(at_new, theta, end):
# create the file name for read in
#---------------------------------
n = float('%.2f' % round(at_new,2))
n = n * 100
na = '%.0f' % round(n)
if len(na) == 3:
na = '0'+na
name = 'astra_{}.{}.001'.format(str(end), na)
x,y,z,px,py,pz,t,q,index,flag = np.loadtxt(name, unpack=True)
rest = np.transpose(np.vstack((t,q,index,flag)))
c = math.cos(theta*math.pi/180)
s = math.sin(theta*math.pi/180)
# positive rotation angle means that new axis are rotated clockwise
m = np.array([[c, 0, s, 0, 0, 0],[0, 1, 0, 0, 0, 0],[-s, 0, c, 0, 0, 0],[0, 0, 0, c, 0, s],[0, 0, 0, 0, 1, 0],[0, 0, 0, -s, 0, c]])
# ASTRA saves x,y,z-z0,px,py,pz-pz0, here add z0 and pz0
z0 = z[0]
pz0 = pz[0]
z[0] = 0
pz[0] = 0
z = z + z0
pz = pz + pz0
prime = np.transpose(np.vstack((x,y,z,px,py,pz)))
# rotation
rot = np.array([m.dot(vec) for vec in prime[:,]])
# subtraction of new z0, x0, pz0. x0 is subtracted in order to be able to use postpro (axis limits in x direction depend on x0). x0 in the new file is not saved!!!
z0 = rot[0,2]
#x0 = apr1(1,1)
pz0 = rot[0,5]
rot[:,2] = rot[:,2]-z0
#apr1(:,1) = apr1(:,1)-x0
rot[:,5] = rot[:,5]-pz0
rot[0,5] = pz0
new = 'rot_' + name
f = open(new,'w')
for i in range(0,len(rot)):
f.write(str('{:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} '.format(rot[i,0],rot[i,1],rot[i,2],rot[i,3],rot[i,4],rot[i,5],t[i],q[i]))+' '+str(int(index[i]))+' '+str(int(flag[i]))+'\n')
f.close()
# rotation yz plane
def rotation_yz(at_new, theta, end):
# create the file name for read in
#---------------------------------
n = float('%.2f' % round(at_new,2))
n = n * 100
na = '%.0f' % round(n)
if len(na) == 3:
na = '0'+na
name = 'astra_{}.{}.001'.format(str(end), na)
x,y,z,px,py,pz,t,q,index,flag = np.loadtxt(name, unpack=True)
rest = np.transpose(np.vstack((t,q,index,flag)))
c = math.cos(theta*math.pi/180)
s = math.sin(theta*math.pi/180)
# positive rotation angle means that new axis are rotated clockwise
m = np.array([[1, 0, 0, 0, 0, 0],[0, c, s, 0, 0, 0],[0, -s, c, 0, 0, 0],[0, 0, 0, 1, 0, 0],[0, 0, 0, 0, c, s],[0, 0, 0, 0, -s, c]])
# ASTRA saves x,y,z-z0,px,py,pz-pz0, here add z0 and pz0
z0 = z[0]
pz0 = pz[0]
z[0] = 0
pz[0] = 0
z = z + z0
pz = pz + pz0
prime = np.transpose(np.vstack((x,y,z,px,py,pz)))
# rotation
rot = np.array([m.dot(vec) for vec in prime[:,]])
# subtraction of new z0, x0, pz0. x0 is subtracted in order to be able to use postpro (axis limits in x direction depend on x0). x0 in the new file is not saved!!!
z0 = rot[0,2]
#x0 = apr1(1,1)
pz0 = rot[0,5]
rot[:,2] = rot[:,2]-z0
#apr1(:,1) = apr1(:,1)-x0
rot[:,5] = rot[:,5]-pz0
rot[0,5] = pz0
new = 'rot_' + name
f = open(new,'w')
for i in range(0,len(rot)):
f.write(str('{:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} '.format(rot[i,0],rot[i,1],rot[i,2],rot[i,3],rot[i,4],rot[i,5],t[i],q[i]))+' '+str(int(index[i]))+' '+str(int(flag[i]))+'\n')
f.close()
# rotation xy plane
def rotation_xy(at_new, theta, end):
# create the file name for read in
#---------------------------------
n = float('%.2f' % round(at_new,2))
n = n * 100
na = '%.0f' % round(n)
if len(na) == 3:
na = '0'+na
name = 'astra_{}.{}.001'.format(str(end), na)
x,y,z,px,py,pz,t,q,index,flag = np.loadtxt(name, unpack=True)
rest = np.transpose(np.vstack((t,q,index,flag)))
c = math.cos(theta*math.pi/180)
s = math.sin(theta*math.pi/180)
# positive rotation angle means that new axis are rotated clockwise
m = np.array([[c, s, 0, 0, 0, 0],[-s, c, 0, 0, 0, 0],[0, 0, 1, 0, 0, 0],[0, 0, 0, c, s, 0],[0, 0, 0, -s, c, 0],[0, 0, 0, 0, 0, 1]])
# ASTRA saves x,y,z-z0,px,py,pz-pz0, here add z0 and pz0
z0 = z[0]
pz0 = pz[0]
z[0] = 0
pz[0] = 0
z = z + z0
pz = pz + pz0
prime = np.transpose(np.vstack((x,y,z,px,py,pz)))
# rotation
rot = np.array([m.dot(vec) for vec in prime[:,]])
# subtraction of new z0, x0, pz0. x0 is subtracted in order to be able to use postpro (axis limits in x direction depend on x0). x0 in the new file is not saved!!!
z0 = rot[0,2]
#x0 = apr1(1,1)
pz0 = rot[0,5]
rot[:,2] = rot[:,2]-z0
#apr1(:,1) = apr1(:,1)-x0
rot[:,5] = rot[:,5]-pz0
rot[0,5] = pz0
new = 'rot_' + name
f = open(new,'w')
for i in range(0,len(rot)):
f.write(str('{:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} {:0.12e} '.format(rot[i,0],rot[i,1],rot[i,2],rot[i,3],rot[i,4],rot[i,5],t[i],q[i]))+' '+str(int(index[i]))+' '+str(int(flag[i]))+'\n')
f.close()
#!/usr/bin/python
import numpy as np
import os
import sys
import argparse
from numpy import sin, cos, pi
from configparser import ConfigParser
acc_dir = os.getenv("ACCDIR")
sys.path.insert(0, os.path.join(acc_dir, "lib"))
from unify import unify, isfloat
from acxml import get_soup
####################################
# Preparation - device description #
####################################
#: creating a file for the solenoid field description
def fieldmap (mapfile):
path, filename = os.path.split(mapfile)
Dir = ('../../elinac/fieldMap/' + filename)
pos, field = np.loadtxt(Dir, usecols=(0, 1), unpack=True)
pos = pos/1000
fp = open('%s' % (filename), 'w')
for k in range(0,len(field)):
fp.write('%0.4f\t%s\n' % (pos[k], field[k]))
fp.close()
return filename
##function called for each know type of element
def SOL (elem, counter, seq_at, sets, element_data):
#Solenoid
counter['sol'] = counter['sol'] + 1
i = str(counter['sol']) #: control variable for counting solenoids
id = elem["id"]
id = id.strip() #: remove whitespaces
at = unify(elem["s"]+"/m") #: location of the solenoid at in m
at = float(at) + seq_at #: location of the solenoid in respect of the beginning of the whole beamline in m
l = float(unify(elem["l"]+"/m")) #: solenoid effective length in m
optr = elem.find("optr")
mapfile = str(optr["mapfile"])
filename = fieldmap(mapfile) #: calls fieldmap function to extract the corresponding fieldmap
scalebz = optr["scalebz"]
ds = 0.1
set_sol = 1
new_at = str(at + l/2 + ds) #: possible end of beamline after ds downstream of the solenoid in m
at = str(at)
l = str(l)
MaxB = str(set_sol/50000) #: scaled Bz (CUR/5/A*G/10000) in Tesla
out_id = str('\n ! '+id)
out = str('\n FILE_BFieLD('+i+')='+filename+', MaxB('+i+')='+MaxB+', S_HIGHER_ORDER('+i+')=T, S_pos('+i+')='+at)
element_data['s_var'].extend(out_id+out)
return element_data,new_at,counter
def MEQ (elem, counter, seq_at, sets, element_data):
#Magnetic or Electrostatic quadrupole
counter['quad'] = counter['quad'] + 1
for se in sets:
if str(elem["id"]+':CUR') == se["pv"]:
quad_val = float(se["value"][:-2])
continue
i = str(counter['quad']) #: control variable for counting quadrupoles
id = elem["id"]
id = id.strip() #: remove whitespaces
at = unify(elem["s"]+"/m") #: location of the quad in the sequence in m
at = float(at) + seq_at #: location of the quad in respect of the beginning of the whole beamline in m
l = float(unify(elem["l"]+"/m")) #: quad effective length in m
optr = elem.find("optr")
aperR = unify(optr["aperr"]+"/m") #: quad apperture/bore in m
btip = optr["btip"]
ds = 0.1 #: look description below (possible end of beamline)
new_at = str(at + l/2 + ds) #: possible end of beamline after ds downstream of the quad in m
at = str(at)
l = str(l)
try:
name,value,unit = btip.split("*")
except:
btip = btip
else:
btip = str(float(value) * quad_val)
out_id = str('\n ! '+id)
out = str('\n Q_length('+i+')='+l+', Q_grad('+i+')='+btip+', Q_bore('+i+')='+aperR+', Q_pos('+i+')='+at)
element_data['q_var'].extend(out_id+out)
return element_data,new_at,counter
def CAV (elem, counter, seq_at, sets, element_data):
#Cavities
counter['cav'] = counter['cav'] + 1
o = 1
for se in sets:
if (elem["id"]+':PHASE' == se["pv"] or elem ["id"][:-3]+'PHASE' or elem["id"][:-4]+'PHASE' == se["pv"]) and o == 1:
phi = se["value"][:-4]
i = str(counter['cav']) #: control variable for counting cavities (including buncher)
id = elem["id"]
id = id.strip() #: remove whitespaces
at = unify(elem["s"]+"/m") #: location of the cavity in m
at = float(at) + seq_at #: location of the cavity in respect of the beginning of the whole beamline in m
l = float(unify(elem["l"]+"/m")) #: cavity length in m
optr = elem.find("optr")
mapfile = str(optr["mapfile"])
freq = unify(optr["freq"]+'/GHz')
ds = 0.1
new_at = str(at + l/2 + ds) #: possible end of beamline after ds downstream of the cavity in m
at = str(at)
l = str(l)
if elem["id"] == 'ELBT:BUNCH':
MaxE = str(0.3)
else:
MaxE = str(10)
out_id = str('\n ! '+id)
out = str('\n FILE_EFieLD('+i+')='+mapfile+', C_HIGHER_ORDER('+i+')=.T, C_SMOOTH('+i+')=5, Nue('+i+')='+freq+', MaxE('+i+')='+MaxE+', Phi('+i+')='+phi+', C_pos('+i+')='+at)
element_data['c_var'].extend(out_id+out)
return element_data,new_at,counter
def BEND (elem, counter, seq_at, sets, element_data):
#Magnetic or Electrostatic dipoles
counter['mb'] = counter['mb'] + 1
for se in sets:
if str(elem["id"]+':CUR') == se["pv"]:
mb_val = float(se["value"][:-2])
i = str(counter['mb']) #: control variable for counting cavities (including buncher)
id = elem["id"]
id = id.strip() #: remove whitespaces
at = float(unify(elem["s"]+"/m")) #: location of the dipole in m
at = float(at) + seq_at #: location of the dipole in respect of the beginning of the whole beamline in m
l = float(unify(elem["l"]+"/m")) #: length of the dipole
w = 0.1 #: in m
ds = 0.3 #: in m
optr = elem.find("optr")
gap = unify(optr["fullgap"]+"/m")
theta = float(unify(optr["theta"]+"/deg")) #: radius of the bend
phi = float(unify(optr["entredge"]+"/deg")) / 180 * pi #: enter and exit edge angle
beta = (90 - float(unify(optr["entredge"]+"/deg"))) / 180 * pi #: 90 deg - enter/exit angle
# ent = unify(optr["entredge"]+"/deg") / 180 * pi #: enter edge angle
# ext = unify(optr["exitedge"]+"/deg") / 180 * pi #: exit edge angle
d11 = str(sin(phi) * w)
d12 = str(at - (sin(phi) * l/2) + (cos(phi) * w))
d21 = str(-(sin(phi) * w))
d22 = str(at - (sin(phi) * l/2) - (cos(phi) * w))
d31 = str(-2 * (cos(phi) * l/2) + (sin(phi) * w))
d32 = str(at + (sin(phi) * l/2) + (cos(phi) * w))
d41 = str(-2 * (cos(phi) * l/2) - (sin(phi) * w))
d42 = str(at + (sin(phi) * l/2) - (cos(phi) * w))
at_new = at + (sin(phi) * l/1) + (cos(beta) * ds)
r = str(l / (2*sin(theta/2)))
at = str(at)
l = str(l)
phi = str(phi)
out_id = str('\n ! '+id)
out = str('\n D_type('+i+')=horizontal, D1('+i+')=('+d11+', '+d12+'), D2('+i+')=('+d21+', '+d22+'), D3('+i+')=('+d31+', '+d32+'), D4('+i+')=('+d41+', '+d42+'), D_radius('+i+')='+r+', D_Gap(1,'+i+')='+gap+', D_Gap(2,'+i+')='+gap+'\n ! run_astra theta = '+str(theta)+'*deg')
element_data['b_var'].extend(out_id+out)
return element_data,at_new,counter
##################################################
# Preparation - device definition and astra file #
##################################################
##Dictionary of elements, read from config.ini file
acc_dir = os.getenv("ACCDIR")
sys.path.insert(0,os.path.join(acc_dir,"lib"))
#: read config file from $ACCDIR/lib/config.ini
config = ConfigParser()
config.read(os.path.join(acc_dir,'lib','config.ini'))
#: create a dictionary of strings (of only the section syf-elements)
elemDict = {s:dict(config.items(s)) for s in config.sections()}['astra-elements']
#: change this dictionary of string into a dictionary of function
for key,value in elemDict.items():
elemDict[key]=locals()[value]
def write_file(stop, element_data, counter, a, pathToOutput):
###set up astra.in header:
header=["&NEWRUN \n",
" Head = 'elinac astra run' \n",
" Run = 1 \n",
" Distribution = 'distribution,ini', Xoff=0.0, Yoff=0.0 \n",
" EmitS = .T \n",
" PhaseS = .T \n",
" TrackS = .T \n",
" RefS = .F \n",
" Qbunch = 8E-3 \n", #bunch charge in nC
" TcheckS = .F \n",
" CathodeS = .T \n",
" TRACK_ALL = .T, PHASE_SCAN = .F, AUTO_PHASE = .F \n",
" check_ref_part = .F \n",
" ZSTART = 0.0, ZSTOP = "+str(stop)+" \n",
" Zemit = 100 \n",
" Zphase = 10 \n",
" H_max = 0.005 \n",
" H_min = 0.001 \n",
"/ \n \n",
"&SCAN \n",
"/ \n \n"
"&CHARGE \n",
" LSPCH = F \n",
" LSPCH3D = F \n",
" Max_scale = 0.05 \n",
" Lmirror = F \n",
" Nxf = 16 \n",
" Nx0 = 2 \n",
" Nyf = 16 \n",
" Ny0 = 2 \n",
" Nzf = 64 \n",
" Nz0 = 2 \n",
" Smooth_x = 0 \n",
" Smooth_y = 0 \n",
" Smooth_z = 0 \n",
"/ \n \n",
"&APERTURE \n",
"/ \n \n"]
##Combine the different pieces of astra.in:
if counter['cav'] == 0:
cavity = [
"&CAVITY \n",
"/ \n \n"]
else:
cavity = ["&CAVITY \n LEFieLD = .T \n"] + element_data['c_var'] + ["/ \n \n"]
if counter['sol'] == 0:
solenoid = [
"&SOLENOID \n",
"/ \n \n"]
else:
solenoid = ["&SOLENOID \n LBFieLD = .T"] + element_data['s_var'] + ["\n/ \n \n"]
if counter['quad'] == 0:
quadrupole = [
"&QUADRUPOLE \n",
"/ \n \n"]
else:
quadrupole = ["&QUADRUPOLE \n LQUAD = .T"] + element_data['q_var'] + ["\n/ \n \n"]
if counter['mb'] == 0:
dipole = [
"&DIPOLE \n",
"/ \n \n"]
else:
dipole = ["&DIPOLE \n LDipole = .T"] + element_data['b_var'] + ["\n/ \n \n"]
asin = header + cavity + solenoid + quadrupole + dipole
a = str(a)
## output to astra.in
if a != 0:
thefile = open(pathToOutput+'astra_'+a+'.in', 'w') # open astra.in to write in it
else:
thefile = open(pathToOutput+'astra_.in', 'w') # open astra.in to write in it
thefile.writelines(asin) # write all lines
thefile.close() # close astra.in
print(a+'. file is converted.')
#########################
# MAIN - call xml2astra #
#########################
def xml2astra(pathToXML, pathToOutput='./', pathToTune=''):
'''conversion from XML to ASTRA'''
soup_path = get_soup(path=pathToXML)
if pathToTune == '':
tunePath = '../tune/'+os.path.basename(pathToXML).split('.',1)[0]+'_ref.xml'
else:
tunePath = pathToTune
soup_tune = get_soup(tune=tunePath)
## initialize variables:
counter = {'cav':0, 'sol':0, 'quad':0, 'mb':0}
a = 0
asin = []
element_data = {'q_var':[], 's_var':[], 'c_var':[], 'b_var':[]}
out = []
out_id = []
at = 0
l = 0
m = 0
#soup = get_soup(beampath=pathToXML)
seqs = soup_path.findAll("seq")
sets = soup_tune.findAll("set")
am = len(seqs)
for seq in seqs:
#reset at-like variable when going to a new sequence
elems = seq.findAll("element")
an = len(elems)
n = 0
for elem in elems:
n = n + 1
#add the element to asin:
typ=elem["type"].replace(' ','')#remove whitespaces. Should also be turned into lower case (how to do that??)
if typ in elemDict:
element_data,at_new,counter = elemDict[typ](elem,counter,at,sets,element_data)
# else:
# print('Element of unknown type: '+typ)
if counter['mb'] == 1:
a = a + 1
write_file(at_new, element_data, counter, a, pathToOutput)
counter = {'cav':0, 'sol':0, 'quad':0, 'mb':0}
at = 0
asin = []
element_data = {'q_var':[], 's_var':[], 'c_var':[], 'b_var':[]}
#rotation_xz(at_new, theta, a)
try: temp = unify(elem["s"]+"/m")
except:
temp = "undefined"
print("Warning in element with id="+id+" attribute 'at' not defined")
if isfloat(temp):
s = unify(elems[-1]["s"]+"/m")
l = unify(elems[-1]["l"]+"/m")
at = at + float(s)
m = m + 1
if n == an and m == am:
if a != 0:
a = a + 1
write_file(at_new, element_data, counter, a, pathToOutput)
print('Well done, your beamline is converted.')
########
# MAIN #
########
#This is just to allow this script to be callable from command line:
if __name__ == '__main__':
#: parsing arguments using library argparse
parser=argparse.ArgumentParser(
description='''Script to convert an acc/*/path XML file into astra input files.''',
epilog="""If you find a bug or want to make a suggestion, create an issue on gitlab.triumf.ca/hla/accsim.""")
parser.add_argument('pathToXML',type=str, default='', help='XML path file name')
parser.add_argument('-t','--tune', type=str, default='', help='XML tune file name')
parser.add_argument('-o','--output', type=str, default='./', help='path where to write astra.in file')
args=parser.parse_args()
xml2astra(args.pathToXML, pathToOutput=args.output, pathToTune=args.tune)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment