<commit_before>"""
This module is a collimator node class for TEAPOT lattice
"""

import os
import math

# import the auxiliary classes
from orbit.utils import orbitFinalize, NamedObject, ParamsDictObject

# import general accelerator elements and lattice
from orbit.lattice import AccNode, AccActionsContainer, AccNodeBunchTracker

# import teapot drift class
from orbit.teapot import DriftTEAPOT

# import Collimator class
from collimator import Collimator

class TeapotCollimatorNode(DriftTEAPOT):
	""" 
	The collimator node class for TEAPOT lattice
	"""
	def __init__(self, length, ma, density_fac, shape, a, b, c, d, angle, name = "collimator no name"):
		"""
		Constructor. Creates the Collimator TEAPOT element.
		"""
		DriftTEAPOT.__init__(self,name)
		self.collimator = Collimator(length,ma,density_fac,shape,a,b,c,d,angle)
		self.setType("collimator teapot")
		self.setLength(length)

	def track(self, paramsDict):
		"""
		The collimator-teapot class implementation of the AccNodeBunchTracker class track(probe) method.
		"""
		length = self.getLength(self.getActivePartIndex())
		bunch = paramsDict["bunch"]
		lostbunch = paramsDict["lostbunch"]
		self.collimator.collimateBunch(bunch, lostbunch)
		#put the track method here
		print "debug tracking the bunch through the collimator name=",self.getName()," part ind=",self.getActivePartIndex()," length=",length
	
	
<commit_msg>Fixed bug.<commit_after>"""
This module is a collimator node class for TEAPOT lattice
"""

import os
import math

# import the auxiliary classes
from orbit.utils import orbitFinalize, NamedObject, ParamsDictObject

# import general accelerator elements and lattice
from orbit.lattice import AccNode, AccActionsContainer, AccNodeBunchTracker

# import teapot drift class
from orbit.teapot import DriftTEAPOT

# import Collimator class
from collimator import Collimator

class TeapotCollimatorNode(DriftTEAPOT):
	""" 
	The collimator node class for TEAPOT lattice
	"""
	def __init__(self, length, ma, density_fac, shape, a, b, c, d, angle, name = "collimator no name"):
		"""
		Constructor. Creates the Collimator TEAPOT element.
		"""
		DriftTEAPOT.__init__(self,name)
		self.collimator = Collimator(length,ma,density_fac,shape,a,b,c,d,angle)
		self.setType("collimator teapot")
		self.setLength(length)

	def track(self, paramsDict):
		"""
		The collimator-teapot class implementation of the AccNodeBunchTracker class track(probe) method.
		"""
		length = self.getLength(self.getActivePartIndex())
		bunch = paramsDict["bunch"]
		lostbunch = paramsDict["lostbunch"]
		self.collimator.collimateBunch(bunch, lostbunch)
		
			

{'fraction_non_alphanumeric': 0.09491006505931879, 'fraction_numerical': 0.0, 'mean_word_length': 5.551378446115288, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '360493', 'n_tokens_mistral': 880, 'n_tokens_neox': 832, 'n_words': 236}