Initial commit

requirements.txt

+numpy>=1.15.4
+pypdt>=0.7

xml2g4bl.py

+#!/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)
+
+
+
+