<commit_before>#!/usr/bin/env python

#--------------------------------------------------------------
# Function to add the aperture class instances to the linac lattice,
# and a function to create the array with aperture nodes and losses.
#--------------------------------------------------------------

import math
import sys
import os

from orbit.py_linac.lattice import LinacApertureNode
from orbit.py_linac.lattice import Quad, Drift, Bend, AbstractRF_Gap
from orbit.py_linac.overlapping_fields import OverlappingQuadsNode

def Add_quad_apertures_to_lattice(accLattice, aprtNodes=[]):
	"""
	Function will add Aperture nodes at the entrance and exit of quads.
	It returns the list of Aperture nodes.
	"""
	node_pos_dict = accLattice.getNodePositionsDict()
	quads = accLattice.getNodesOfClasses([Quad,OverlappingQuadsNode])
	for node in quads:
		if(isinstance(node,Quad)):
			if(node.hasParam("aperture") and node.hasParam("aprt_type")):
				shape = node.getParam("aprt_type")
				a = node.getParam("aperture")
				node_name = node.getName()
				(posBefore, posAfter) = node_pos_dict[node]
				apertureNodeBefore = LinacApertureNode(shape,a/2.0,a/2.0,posBefore)
				apertureNodeAfter = LinacApertureNode(shape,a/2.0,a/2.0,posAfter)
				apertureNodeBefore.setName(node_name+":AprtIn")
				apertureNodeAfter.setName(node_name+":AprtOut")
				apertureNodeBefore.setSequence(node.getSequence())
				apertureNodeAfter.setSequence(node.getSequence())
				node.addChildNode(apertureNodeBefore,node.ENTRANCE)
				node.addChildNode(apertureNodeAfter,node.EXIT)
				aprtNodes.append(apertureNodeBefore)
				aprtNodes.append(apertureNodeAfter)
		if(isinstance(node,OverlappingQuadsNode)):
			# place to add aperture for the overlapping fields quads
			nParts = node.getnParts()
			simple_quads = node.getQuads()
			quad_centers = node.getCentersOfField()
			node_start_pos = node_pos_dict[node][0]
			pos_part_arr = []
			s = 0.
			for part_ind in range(nParts):
				pos_part_arr.append(s)
				s += node.getLength(part_ind)
			for quad_ind in range(len(simple_quads)):
				quad = simple_quads[quad_ind]
				shape = quad.getParam("aprt_type")
				a = quad.getParam("aperture")
				quad_name = quad.getName()				
				length = quad.getLength()
				pos_z = quad_centers[quad_ind]
				posBefore = pos_z - length/2.
				posAfter = pos_z + length/2.
				for part_ind in range(nParts-1):
					if(posBefore >= pos_part_arr[part_ind] and posBefore < pos_part_arr[part_ind+1]):
						apertureNodeBefore = LinacApertureNode(shape,a/2.0,a/2.0,posBefore + node_start_pos)
						apertureNodeBefore.setName(quad_name+":AprtIn")
						apertureNodeBefore.setSequence(node.getSequence())
						node.addChildNode(apertureNodeBefore,node.BODY,part_ind,node.BEFORE)
						aprtNodes.append(apertureNodeBefore)
					if(posAfter > pos_part_arr[part_ind] and posAfter <= pos_part_arr[part_ind+1]):
						apertureNodeAfter = LinacApertureNode(shape,a/2.0,a/2.0,posAfter + node_start_pos)
						apertureNodeAfter.setName(quad_name+":AprtOut")
						apertureNodeAfter.setSequence(node.getSequence())
						node.addChildNode(apertureNodeAfter,node.BODY,part_ind,node.AFTER)
						aprtNodes.append(apertureNodeAfter)
	return aprtNodes
	

def Add_bend_apertures_to_lattice(accLattice, aprtNodes=[], step = 0.25):
	"""
	Function will add Aperture nodes at the bend nodes of the lattice if they
	have aperture parameters.
	The distance between aperture nodes is defined by the step input parameter 
	(in meters).
	"""
	node_pos_dict = accLattice.getNodePositionsDict()
	bends = accLattice.getNodesOfClasses([Bend,])
	if(len(bends) <= 0): return aprtNodes
	for node in bends:
		if(node.hasParam("aperture_x") and node.hasParam("aperture_y") and node.hasParam("aprt_type")): 
			shape = node.getParam("aprt_type")
			a_x = node.getParam("aperture_x")			
			a_y = node.getParam("aperture_y")
			if(shape != 3): continue
			node_name = node.getName()
			(posBefore, posAfter) = node_pos_dict[node]
			nParts = node.getnParts()
			s = posBefore
			pos_part_arr = []
			for part_ind in range(nParts):
				pos_part_arr.append(s)
				s += node.getLength(part_ind)
			run_path = pos_part_arr[0]
			run_path_old = pos_part_arr[0] - 2*step
			for part_ind in range(nParts):
				run_path = pos_part_arr[part_ind]
				if(run_path >= run_path_old + step):	
					apertureNode = LinacApertureNode(shape,a_x/2.0,a_y/2.0,run_path)
					aprt_name = node_name+":"+str(part_ind)+":Aprt"
					if(nParts == 1): aprt_name = node_name+":Aprt"					
					apertureNode.setName(aprt_name)
					apertureNode.setSequence(node.getSequence())
					node.addChildNode(apertureNode,node.BODY,part_ind,node.BEFORE)		
					aprtNodes.append(apertureNode)
					run_path_old = run_path
	return aprtNodes			


def Add_rfgap_apertures_to_lattice(accLattice, aprtNodes=[]):
	"""
	Function will add Aperture nodes at the entrance and exit of RF gap.
	It returns the list of Aperture nodes.
	"""
	node_pos_dict = accLattice.getNodePositionsDict()
	rfgaps = accLattice.getNodesOfClasses([AbstractRF_Gap,])
	for node in rfgaps:
		if(node.hasParam("aperture") and node.hasParam("aprt_type")):
			shape = node.getParam("aprt_type")
			a = node.getParam("aperture")
			node_name = node.getName()
			(posBefore, posAfter) = node_pos_dict[node]<commit_msg>At this moment the linac RF gap apertures can be added only for BaseRF_Gap. It will change in the future. We plan to add another subclass of the AbstractRF_Gap.<commit_after>#!/usr/bin/env python

#--------------------------------------------------------------
# Function to add the aperture class instances to the linac lattice,
# and a function to create the array with aperture nodes and losses.
#--------------------------------------------------------------

import math
import sys
import os

from orbit.py_linac.lattice import LinacApertureNode
from orbit.py_linac.lattice import Quad, Drift, Bend, BaseRF_Gap
from orbit.py_linac.overlapping_fields import OverlappingQuadsNode

def Add_quad_apertures_to_lattice(accLattice, aprtNodes=[]):
	"""
	Function will add Aperture nodes at the entrance and exit of quads.
	It returns the list of Aperture nodes.
	"""
	node_pos_dict = accLattice.getNodePositionsDict()
	quads = accLattice.getNodesOfClasses([Quad,OverlappingQuadsNode])
	for node in quads:
		if(isinstance(node,Quad)):
			if(node.hasParam("aperture") and node.hasParam("aprt_type")):
				shape = node.getParam("aprt_type")
				a = node.getParam("aperture")
				node_name = node.getName()
				(posBefore, posAfter) = node_pos_dict[node]
				apertureNodeBefore = LinacApertureNode(shape,a/2.0,a/2.0,posBefore)
				apertureNodeAfter = LinacApertureNode(shape,a/2.0,a/2.0,posAfter)
				apertureNodeBefore.setName(node_name+":AprtIn")
				apertureNodeAfter.setName(node_name+":AprtOut")
				apertureNodeBefore.setSequence(node.getSequence())
				apertureNodeAfter.setSequence(node.getSequence())
				node.addChildNode(apertureNodeBefore,node.ENTRANCE)
				node.addChildNode(apertureNodeAfter,node.EXIT)
				aprtNodes.append(apertureNodeBefore)
				aprtNodes.append(apertureNodeAfter)
		if(isinstance(node,OverlappingQuadsNode)):
			# place to add aperture for the overlapping fields quads
			nParts = node.getnParts()
			simple_quads = node.getQuads()
			quad_centers = node.getCentersOfField()
			node_start_pos = node_pos_dict[node][0]
			pos_part_arr = []
			s = 0.
			for part_ind in range(nParts):
				pos_part_arr.append(s)
				s += node.getLength(part_ind)
			for quad_ind in range(len(simple_quads)):
				quad = simple_quads[quad_ind]
				shape = quad.getParam("aprt_type")
				a = quad.getParam("aperture")
				quad_name = quad.getName()				
				length = quad.getLength()
				pos_z = quad_centers[quad_ind]
				posBefore = pos_z - length/2.
				posAfter = pos_z + length/2.
				for part_ind in range(nParts-1):
					if(posBefore >= pos_part_arr[part_ind] and posBefore < pos_part_arr[part_ind+1]):
						apertureNodeBefore = LinacApertureNode(shape,a/2.0,a/2.0,posBefore + node_start_pos)
						apertureNodeBefore.setName(quad_name+":AprtIn")
						apertureNodeBefore.setSequence(node.getSequence())
						node.addChildNode(apertureNodeBefore,node.BODY,part_ind,node.BEFORE)
						aprtNodes.append(apertureNodeBefore)
					if(posAfter > pos_part_arr[part_ind] and posAfter <= pos_part_arr[part_ind+1]):
						apertureNodeAfter = LinacApertureNode(shape,a/2.0,a/2.0,posAfter + node_start_pos)
						apertureNodeAfter.setName(quad_name+":AprtOut")
						apertureNodeAfter.setSequence(node.getSequence())
						node.addChildNode(apertureNodeAfter,node.BODY,part_ind,node.AFTER)
						aprtNodes.append(apertureNodeAfter)
	return aprtNodes
	

def Add_bend_apertures_to_lattice(accLattice, aprtNodes=[], step = 0.25):
	"""
	Function will add Aperture nodes at the bend nodes of the lattice if they
	have aperture parameters.
	The distance between aperture nodes is defined by the step input parameter 
	(in meters).
	"""
	node_pos_dict = accLattice.getNodePositionsDict()
	bends = accLattice.getNodesOfClasses([Bend,])
	if(len(bends) <= 0): return aprtNodes
	for node in bends:
		if(node.hasParam("aperture_x") and node.hasParam("aperture_y") and node.hasParam("aprt_type")): 
			shape = node.getParam("aprt_type")
			a_x = node.getParam("aperture_x")			
			a_y = node.getParam("aperture_y")
			if(shape != 3): continue
			node_name = node.getName()
			(posBefore, posAfter) = node_pos_dict[node]
			nParts = node.getnParts()
			s = posBefore
			pos_part_arr = []
			for part_ind in range(nParts):
				pos_part_arr.append(s)
				s += node.getLength(part_ind)
			run_path = pos_part_arr[0]
			run_path_old = pos_part_arr[0] - 2*step
			for part_ind in range(nParts):
				run_path = pos_part_arr[part_ind]
				if(run_path >= run_path_old + step):	
					apertureNode = LinacApertureNode(shape,a_x/2.0,a_y/2.0,run_path)
					aprt_name = node_name+":"+str(part_ind)+":Aprt"
					if(nParts == 1): aprt_name = node_name+":Aprt"					
					apertureNode.setName(aprt_name)
					apertureNode.setSequence(node.getSequence())
					node.addChildNode(apertureNode,node.BODY,part_ind,node.BEFORE)		
					aprtNodes.append(apertureNode)
					run_path_old = run_path
	return aprtNodes			


def Add_rfgap_apertures_to_lattice(accLattice, aprtNodes=[]):
	"""
	Function will add Aperture nodes at the entrance and exit of RF gap.
	It returns the list of Aperture nodes.
	"""
	node_pos_dict = accLattice.getNodePositionsDict()
	rfgaps = accLattice.getNodesOfClasses([BaseRF_Gap,])
	for node in rfgaps:
		if(node.hasParam("aperture") and node.hasParam("aprt_type")):
			shape = node.getParam("aprt_type")
			a = node.getParam("aperture")
			node_name = node.getName()
			(posBefore, posAfter) = node_pos_dict[node]
{'fraction_non_alphanumeric': 0.13085076446862395, 'fraction_numerical': 0.0067535878435418814, 'mean_word_length': 4.165697674418604, 'pattern_counts': {'":': 14, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5629011', 'n_tokens_mistral': 4344, 'n_tokens_neox': 3801, 'n_words': 693}