Commit 89ed9117 authored by Paul Jung's avatar Paul Jung
Browse files

Initial commit

parents
Pipeline #5687 failed with stages
#!/usr/bin/python
# Known Problems to date:
# 1) Sequence resets the "at", need to track "at"s as you go through sequences - used last_at to get around for now
# 2) Need beam pipe throughout line
# 3) Corner correction - from genericbend
# 4) Gradients should be coming from tune
# 5) WF field should come from optr radius, E from v*B
import numpy as np
import pypdt
from math import sqrt
import re
import os
import sys
from configparser import ConfigParser
import argparse
acc_dir = os.getenv("ACCDIR")
sys.path.insert(0, os.path.join(acc_dir, "lib"))
from unify import unify, isfloat
from scale import load_beam_properties
from acxml import get_soup, get_root
## Element Functions
#def FIT(elem):
# # Will do something with this and GMINUIT if I find the time
# return '',''
def MQ(elem,last_at):
# Magnetic Quad -- genericquad
id=elem["id"]
id=id.replace(' ','')#remove whitespaces
at="%.2f" % (float(unify(elem["s"]+"/mm"))+last_at) #Location of quad at center in mm
g4bl=elem.find("g4bl")
optr=elem.find("optr")
l=unify(elem["l"]+"/mm")
if float(l)==0:
l=unify(g4bl["fieldl"]+"/mm")
ironR=unify(g4bl["ironr"]+"/mm")
aperR="%.3f" % abs(float(unify(optr["aperr"]+"/mm")))
coilR=unify(g4bl["coilr"]+"/mm")
coilHW=unify(g4bl["coilhw"]+"/mm")
try:
ironL=g4bl["ironl"]
except:
ironL=str(float(l)+50)
grad=unify(g4bl["grad"]+"/T") # Gradient should come from tune, do this just for proof of concept
try:
ironM=g4bl["ironm"]
except:
ironM="Fe"
try:
fieldM=g4bl["fieldm"]
except:
fieldM="Vacuum"
# Should add a kill paramater to build command, will likely require a g4bl element
build_qf="tubs "+id+"_QF"+" length=0.5 color=1,0.5,0,0.2 kill=1"
build_cmd="genericquad "+id+" fieldLength="+l+" ironLength="+ironL+" ironRadius="+ironR+ \
" poleTipRadius="+aperR+" coilRadius="+coilR+" coilHalfWidth="+coilHW+ \
" gradient="+grad+" ironMaterial="+ironM+" fieldMaterial="+fieldM+" ironColor=0.5,0.1,0.1,0.3"+ \
" fringe=0 kill=1"
place_qf1="place "+id+"_QF z="+"%.3f" % (float(at)-float(ironL)/2-0.25)+" innerRadius="+"%.3f" % (float(aperR)*0.99)+ " outerRadius="+ironR
place_cmd="place "+id+" z="+at
place_qf2="place "+id+"_QF z="+"%.3f" % (float(at)+float(ironL)/2+0.25)+" innerRadius="+"%.3f" % (float(aperR)*0.99)+ " outerRadius="+ironR
build_out=["# Magnetic Quadrupole "+id,
build_cmd,"",
"# Quad Fittings to approx beampipe inside quad "+id,
build_qf,""]
place_out=[place_qf1,place_cmd,place_qf2]
return build_out,place_out
def MB(elem,last_at):
# Magnetic Bending Dipole -- genericbend
id=elem["id"]
id=id.replace(' ','')
at="%.2f" % (float(unify(elem["s"]+"/mm"))+last_at)
l=unify(elem["l"]+"/mm")
g4bl=elem.find("g4bl")
optr=elem.find("optr")
theta=unify(optr["theta"]+"/deg")
entrEdge=unify(optr["entredge"]+"/deg")
exitEdge=unify(optr["exitedge"]+"/deg")
iron_dim, field_dim, (fieldM,ironM)=elem_size(elem) # Cleaner than putting try's here
ironW,ironH,ironL=iron_dim[0],iron_dim[1],iron_dim[2]
fieldW,fieldH,l=field_dim[0],field_dim[1],field_dim[2]
if isfloat(theta) and isfloat(l):
lf=float(l)
thf=float(theta)
if thf != 0.0:
r=lf/(thf*np.pi/180.0)
R="%.3f"%r
ByF=unify(g4bl["byf"])
build_cmd="genericbend "+id+" By="+ByF+" fieldWidth="+fieldW+" fieldHeight="+fieldH+" fieldLength="+l+" ironMaterial="+ironM+\
" fieldMaterial="+fieldM+" ironWidth="+ironW+" ironHeight="+ironH+" ironLength="+ironL+" ironColor=0.7,0.7,0.7,0.4 kill=1"
place_mb="place "+id+" z="+at+" rotation=Y-"+entrEdge
place_corn="corner C_"+id+" z="+at+" rotation=Y-"+theta+" radiusCut=500"
build_out=["# Magnetic Bending Dipole "+id,
build_cmd,""]
place_out=[place_mb,place_corn]
return build_out,place_out
def WF(elem,last_at):
# Wien Filter -- Custom built from box, approximation, need fringe fields
id=elem["id"]
id=id.replace(' ','')
at="%.2f" % (float(unify(elem["s"]+"/mm"))+last_at)
l=unify(elem["l"]+"/mm")
g4bl=elem.find("g4bl")
# Get iron and field size
iron_dim, field_dim, (fieldM,ironM)=elem_size(elem) # Cleaner than putting try's here
ironW,ironH,ironL=iron_dim[0],iron_dim[1],iron_dim[2]
fieldW,fieldH,l=field_dim[0],field_dim[1],field_dim[2]
WFBy=unify(g4bl["wfby"]+"/T")
build_cmd="box "+id+"Outer height="+ironH+" width="+ironW+" length="+ironL+" material="+ironM+" color=0.2,0.2,1,0.5 kill=1 \n"+\
"box "+id+"Inner height="+fieldH+" width="+fieldW+" length="+l+" material="+fieldM+" color=0.2,0.6,0.4,0.5 kill=0 \n"+\
"fieldexpr "+id+"Field By="+WFBy+" Ex=80.6042*"+WFBy+" length="+l+" width="+fieldW+" height="+fieldH+" \n"
group_cmd="group "+id+" \n"\
"\t place "+id+"Field parent="+id+"Inner \n"+\
"\t place "+id+"Inner parent="+id+"Outer \n"+\
"\t place "+id+"Outer \n"+\
"endgroup"
place_cmd="place "+id+" z="+at
build_out=["# Wien Filter "+id,
build_cmd,
"# Now group the elements into a WF",
group_cmd,""]
place_out=[place_cmd]
return build_out,place_out
def MARKER(elem,last_at):
# User Marker -- Use 'id' as comment
id=elem["id"]
id=id.replace(' ','')#remove whitespace
at=u"%.2f" % (float(unify(elem["s"]+"/mm"))+last_at)
return ["# Build End: "+id+"\n"],['# Place End: '+id+"\n"]
## Script Functions
#: split_str - Cuts strings for clean output, adds '\' for shell-style line continuation
#: elem_size - return various parameters that could or could not be tagged
#: formatg4bl - Adds new lines between strings for output, splits if lines are too long
#: elem2g4bl - Returns the element
#: xml2g4bl - Converts the XML file to .g4bl file
def split_str(string,size):
## Function to split strings into lengths of [size] for clean output
## Called in formatg4bl iteratively
cut=string[:size]
out_string='\t'+string[cut.rfind(' '):]
cut=cut[:cut.rfind(' ')]+' \\\n'
return cut, out_string
def elem_size(elem):
# Function for getting the various element parameters that I "try" to get
# Probably a cleaner way of doing this, but I wanted something quick and easy
g4bl=elem.find("g4bl")
fieldL=unify(elem["l"]+"/mm")
fieldW=unify(g4bl["fieldw"]+"/mm")
fieldH=unify(g4bl["fieldh"]+"/mm")
# If iron parameters given, use them, otherwise default
try:
ironW=unify(g4bl["ironw"]+"/mm")
except:
ironW="%.3f" % (float(fieldW)+100)
try:
ironH=unify(g4bl["ironh"]+"/mm")
except:
ironH="%.3f" % (float(fieldH)+100)
try:
ironL=unify(g4bl["ironl"]+"/mm")
except:
ironL=unify(elem["l"]+"/mm")
# Now get materials, default to iron and vacuum
try:
ironM=g4bl["ironm"]
except:
ironM="Fe"
try:
fieldM=g4bl["fieldm"]
except:
fieldM="Vacuum"
return (ironW,ironH,ironL), (fieldW,fieldH,fieldL), (fieldM,ironM)
def formatg4bl(g4bl,size):
## Format output of g4bl string to be clean for an input file
## Constrains line length to [size] and adds '\' for end of split lines
## Basically, add a newline in between lines, split recursively if len > size
line_lst=[]
for line in g4bl:
if len(line)>=size:
spl_line=line.split('\n')
spl_gr=False
for spl in spl_line:
if len(spl) >= size:
spl_gr=True
if spl_gr==True:
# Still need to tidy this up loop up, but it is functional for now
while len(line) >= size:
cut,line=split_str(line,size)
line_lst.append(cut)
line_lst.append(line+'\n')
# If spl_gr is false, just write the line, it'll split itself
elif spl_gr==False:
line_lst.append(line+'\n')
else:
line=line+'\n'
line_lst.append(line)
return line_lst
# Could use xml2syf.elem2optr(elem), not sure if being clear is better than concise
# ^ -- Slightly modified version, might as well just use this
def elem2g4bl(elem,last_at):
type=elem["type"].replace(' ','')
if type in elemDict:
return elemDict[type](elem,last_at)[0],elemDict[type](elem,last_at)[1]
else:
print('Element of unknown type: '+type)
out_line=['#Element of unknown type: '+type+" at s="+unify(elem["s"]+"/mm")]
return out_line,out_line
def xml2g4bl(pathToXML,pathToTuneFile='',pathToSigmaFile='',startId='',endId='',pathToOutput='./'):
'''Conversion from XML to G4BL'''
#: Get XML soup
soup_path=get_soup(path=pathToXML)
##Initializa variables
#: Write to .g4bl only when flag==1
flag=1
if startId != '':
flag=0
g4bl_build=[]
g4bl_place=[]
at=0
l=0
prev_at=0
last_at=0
prev_l=0
##Parse XML and start building g4bl (to be written as pathToXML.g4bl)
seqs = soup_path.findAll("seq")
for seq in seqs:
#reset at-like variable when going to a new sequence
drift_l=0.0
prev_at=0.0
prev_l=0.0
elems=seq.findAll("element")
for elem in elems:
drift_l=0.0
try: id=elem["id"]
except:
id="undefined"
print(["Warning: id not defined; element=",elem])
try: temp=unify(elem["s"]+"/mm")
except:
temp="undefined"
at=prev_at+last_at
print("Warning in element with id="+id+" attribute 'at' not defined")
if isfloat(temp):
at=float(temp)+last_at
try: temp=unify(elem["l"]+"/mm")
except:
temp="undefined"
l=0.0
print("Warning in element with id="+id+" attribute 'l' not defeind")
if isfloat(temp):
l=float(temp)
else:
l=0.0 # if l not specified, it is assumed to be zero-length
#drift length b/w this element and the previous one:
drift_l=at-prev_at-l/2.0-prev_l/2.0+last_at
else:
at=prev_at+last_at
prev_l=l
prev_at=at
if startId != '':
if id==startId:
flag=1
#: use a build and place string to keep things organized
#: ^ -- Final file will call build first, then place
if flag==1:
g4bl_build.extend(elem2g4bl(elem,last_at)[0])
g4bl_place.extend(elem2g4bl(elem,last_at)[1])
#: to handle stopping at element with id=endId
if endId != '':
if id==endId:
flag=0
last_at=at
# Get Sigma info -> For beam paramaters
try:
#default sigma file name is the same than the XML path file name:
if pathToSigmaFile == '':
pathToSigmaFile=pathToXML
soup_sigma = get_soup(sigma=pathToSigmaFile)
#if sigma file not found, use the default one:
except:
#get the accelerator (=isac, or elinacm etc) from soup_path (was added there AUTOMATICALLY by get_soup)
acc=soup_path.find('body')['acc']
#use the default sigma XML for this particular accelerator:
try:
soup_sigma = get_soup(sigma='default.xml',acc=acc)
print('sigma XML file '+pathToSigmaFile+' not found. Using the default sigma file instead.')
except:
print("Error: no sigma file found in $ACCDIR/"+acc)
exit(1)
temp=soup_sigma.find('beampara').get_text()
sigma=temp.split('\n')
# if first space is zero, we want sigma[1] for beampara, remove leading whitespace
if (sigma[0].isspace() or sigma[0]==''):
para_line=sigma[1].lstrip()
unit_line=sigma[2].lstrip()
else:
para_line=sigma[0].lstrip()
unit_line=sigma[1].lstrip()
# Remove comments, make string array with numeric values as entries
para_line=para_line.split('!')[0]
unit_line=unit_line.split('!')[0]
para_line=para_line.split()
unit_line=unit_line.split()
# Get momentum value from tune, for beam command p=sqrt[(K+m)^2-m^2]
tunePath=pathToTuneFile if pathToTuneFile else os.path.basename(pathToXML).split('.',1)[0]+'_ref.xml'
soup_tune=get_root(tune=tunePath)
top_tune=soup_tune.find('tune')
[ID,massMeV,chargeState,EkMeV]=load_beam_properties(top_tune)
massMeV=unify(massMeV+'/MeV')
EkMeV=unify(EkMeV+'/MeV')
charge=float(chargeState)
# Determine the particle based on the particle mass by comparing to masses in PyPDT
# Probably not the best way, but a quick solution
tbl=pypdt.ParticleDataTable()
try:
for p in tbl:
if (abs(float(massMeV)- 1000*p.mass)) <=0.001:
pid=p.name
pid=re.sub("[^a-zA-Z]+", "", pid)
print(charge/abs(charge))
if (charge/abs(charge) == 1):
pid=pid+"+"
elif (charge/abs(charge) == -1):
pid=pid+"-"
else:
print("Could not determine sign of particle charge")
except:
print("Error, could not match massMeV to particle in PyPDT")
# Get values from sigma, for beam command -> The x',y',delta don't actually line up perfectly, but give a close start
beamname=pathToXML
# pid="mu+" #get from mass -> lookup table/dict?
meanP=str(sqrt((float(EkMeV)+float(massMeV))*(float(EkMeV)+float(massMeV))-float(massMeV)*float(massMeV)))
sigmaP=para_line[5]
sigmaX=unify(para_line[0]+"*"+unit_line[0]+"*cm/mm")
sigmaY=unify(para_line[2]+"*"+unit_line[2]+"*cm/mm")
sigmaXp=para_line[1]
sigmaYp=para_line[3]
## Build strings for input file, use -unset for things people may want to set from command line
header=["# G4beamline file for "+beamname,
"",
"physics QGSP_BERT spinTracking=1",
"param worldMaterial=Vacuum",
"param fieldVoxel=20,20,20",
"param eventTimeLimit=1",
"param maxStep=5",
"param -unset nEvents=10000",
"trackcolor reference=1,1,1",
""]
beam_cmd=["# Beam Parameters",
"param -unset meanP="+meanP,
"param -unset sigmaP="+sigmaP,
"param -unset sigmaX="+sigmaX,
"param -unset sigmaY="+sigmaY,
"param -unset sigmaXp="+sigmaXp,
"param -unset sigmaYp="+sigmaYp,
"",
"beam ellipse particle="+pid+" nEvents=$nEvents meanMomentum=$meanP sigmaP=$sigmaP "+
"sigmaX=$sigmaX sigmaY=$sigmaY sigmaXp=$sigmaXp sigmaYp=$sigmaYp beamZ=0","",
"reference particle="+pid+" referenceMomentum=$meanP noEloss=1 noEfield=0 beamX=0 beamY=0 beamZ=0",""]
footer=[""]
# Header: General Simulation parameters
# Beam Command: The beam parameters, and command for beam/reference
# g4bl_build: The commands that actually build elements
# g4bl_place: The commands that place the elements in the beamline
# Footer: Put output stuff here? Or user defined output?
g4bl=header+beam_cmd+["## Now Build the Elements"]+g4bl_build+["## Now Place the Elements"]+g4bl_place+footer
g4bl=formatg4bl(g4bl,80) # Format g4bl string as defined in the formatg4bl() func
thefile = open(pathToOutput+beamname+'.g4bl','w') # Open beamname.g4bl to write in it
thefile.writelines(g4bl)
thefile.close() #close g4bl file
##Dictionary of elements, read from config.ini file
dir_path = os.path.dirname(os.path.realpath(__file__))
Config = ConfigParser()
#: Read config file - I've been using this configg4bl.ini
Config.read(dir_path+"/config.ini")
#: Create a dictionary of strings (of only the syf-elements)
elemDict = {s:dict(Config.items(s)) for s in Config.sections()}['g4bl-elements']
#: Change the dictionary of strings into a dictionary of functions, this writes out file
for key,value in elemDict.items():
elemDict[key]=locals()[value]
#This allows the script to be called from command line
if __name__ == '__main__':
#: parsing arguments using library argparse
parser=argparse.ArgumentParser(
description='''Script to convert acc/*/path XML file into g4bl input files.''',
epilog='''If you find a bu or want to make a suggestion, create an issue on gitlab.triumf.ca/hla/acc.''')
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('-s','--sigma', type=str, default='', help='XML sigma file name')
parser.add_argument('-st','--start', type=str, default='', help='Provide an element id from which you want the conversion to start')
parser.add_argument('-ed','--end', type=str, default='', help='Provide an element id at which you want the conversion to end')
parser.add_argument('-o','--output', type=str, default='./', help='path where to write sy.f and data.dat files')
args=parser.parse_args()
xml2g4bl(args.pathToXML, pathToTuneFile=args.tune, pathToSigmaFile=args.sigma,pathToOutput=args.output, startId=args.start, endId=args.end)
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