diff --git a/py/orbit/foils/TeapotFoilNode.py b/py/orbit/foils/TeapotFoilNode.py index b8b7a797..f37448de 100644 --- a/py/orbit/foils/TeapotFoilNode.py +++ b/py/orbit/foils/TeapotFoilNode.py @@ -48,3 +48,7 @@ def track(self, paramsDict): def setScatterChoice(self, choice): self.scatterChoice = choice + + def getnHits(self): + #returns the particle foil hits + return self.foil.getnHits() \ No newline at end of file diff --git a/src/orbit/MaterialInteractions/Foil.cc b/src/orbit/MaterialInteractions/Foil.cc index a0eaa2b9..9da7b2f6 100644 --- a/src/orbit/MaterialInteractions/Foil.cc +++ b/src/orbit/MaterialInteractions/Foil.cc @@ -44,6 +44,8 @@ Foil::Foil(double xmin, double xmax, double ymin, double ymax, double thick): Cp thick_ = thick; length_ = 0.0; ma_ = 0; + nHits = 0; + nLost = 0; } /////////////////////////////////////////////////////////////////////////// @@ -606,3 +608,19 @@ void Foil::loseParticle(Bunch* bunch, Bunch* lostbunch, int ip, int& nLost, int& zrl = -1.; } + +/////////////////////////////////////////////////////////////////////////// +// +// NAME +// +// Foil::getnHits +// +// DESCRIPTION +// +// Returns the accumulated number of foil hits. +// +/////////////////////////////////////////////////////////////////////////// + +int Foil::getnHits(){ + return nHits; +} diff --git a/src/orbit/MaterialInteractions/Foil.hh b/src/orbit/MaterialInteractions/Foil.hh index 3b2052ae..5e6bd260 100644 --- a/src/orbit/MaterialInteractions/Foil.hh +++ b/src/orbit/MaterialInteractions/Foil.hh @@ -25,6 +25,9 @@ public: /** Routine for transfering particles through a foil with simplified scattering. No particle loss. */ void traverseFoilSimpleScatter(Bunch* bunch); + /** Returns the accumulated number of foil hits. */ + int getnHits(); + private: diff --git a/src/orbit/MaterialInteractions/wrap_foil.cc b/src/orbit/MaterialInteractions/wrap_foil.cc index 3d9398a9..8e10e7ba 100644 --- a/src/orbit/MaterialInteractions/wrap_foil.cc +++ b/src/orbit/MaterialInteractions/wrap_foil.cc @@ -80,6 +80,12 @@ extern "C" { return Py_None; } + /** Returns the accumulated number of foil hits. */ + static PyObject* Foil_getnHits(PyObject *self, PyObject *args){ + Foil* cpp_Foil = (Foil*)((pyORBIT_Object*) self)->cpp_obj; + return PyLong_FromLong(cpp_Foil->getnHits()); + } + //----------------------------------------------------- @@ -96,6 +102,7 @@ extern "C" { static PyMethodDef FoilClassMethods[] = { { "traverseFoilFullScatter", Foil_traverseFoilFullScatter, METH_VARARGS,"Performs the foil scatter of the bunch."}, { "traverseFoilSimpleScatter", Foil_traverseFoilSimpleScatter, METH_VARARGS,"Performs the foil scatter of the bunch."}, + {"getnHits", Foil_getnHits, METH_NOARGS, "Returns the accumulated number of foil hits. "}, {NULL} }; diff --git a/tests/examples/Foil/test_Foil_getnHits.py b/tests/examples/Foil/test_Foil_getnHits.py new file mode 100644 index 00000000..781a7ada --- /dev/null +++ b/tests/examples/Foil/test_Foil_getnHits.py @@ -0,0 +1,156 @@ +############################################################## +# This script reads the input MAD file with lattice information, +# creates the TEAPOT lattice, and modifies this lattice by inserting +# several collimation nodes +############################################################## +import pytest +import os +import random + +from orbit.core.orbit_utils import random as orbit_random +from orbit.teapot import teapot +from orbit.teapot import DriftTEAPOT +from orbit.core.bunch import Bunch +from orbit.utils.orbit_mpi_utils import bunch_orbit_to_pyorbit, bunch_pyorbit_to_orbit +from orbit.foils import TeapotFoilNode +from orbit.foils import addTeapotFoilNode + + +def read_values_from_file(file_path): + values = [] + + with open(file_path) as f: + for line in f: + line = line.strip() + if line: + values.extend(map(float, line.split())) + + return values + + +random.seed(100) +orbit_random.seed(0) +script_dir = os.path.dirname(__file__) +print("Start.") + +teapot_latt = teapot.TEAPOT_Lattice() +print("Read MAD.") +teapot_latt.readMAD(os.path.join(script_dir, "../Apertures/MAD_Lattice/LATTICE"), "RING") +print("Lattice=", teapot_latt.getName(), " length [m] =", teapot_latt.getLength(), " nodes=", len(teapot_latt.getNodes())) + +n_drifts = 0 +drift_length = 0.0 +for node in teapot_latt.getNodes(): + if isinstance(node, DriftTEAPOT): + n_drifts += 1 + drift_length += node.getLength() + +print("number of drifts =", n_drifts) +print("total drift length =", drift_length) +print("============= nodes inside the region ===========") +pos_start = 16.5 +pos_stop = 20.5 +# print all nodes around the specified position +for node in teapot_latt.getNodes(): + (node_pos_start, node_pos_stop) = teapot_latt.getNodePositionsDict()[node] + if node_pos_start > pos_start and node_pos_stop < pos_stop: + print("node=", node.getName(), " type=", node.getType(), " pos=", node_pos_start, " L=", node.getLength()) + + +xmin = -0.100 +xmax = 0.100 +ymin = -0.100 +ymax = 0.100 +# Below is 1000 times the width of normal foil but will do only one turn. +thick = 40000 + +foil = TeapotFoilNode(xmin, xmax, ymin, ymax, thick, "Foil 1") + +addTeapotFoilNode(teapot_latt, 18.5, foil) +print("===========Lattice modified =======================================") +print("New Lattice=", teapot_latt.getName(), " length [m] =", teapot_latt.getLength(), " nodes=", len(teapot_latt.getNodes())) + +n_drifts = 0 +drift_length = 0.0 +for node in teapot_latt.getNodes(): + if isinstance(node, DriftTEAPOT): + n_drifts += 1 + drift_length += node.getLength() + +print("number of drifts =", n_drifts) +print("total drift length =", drift_length) + +print("============= nodes inside the region ===========") +# print all nodes around the specified position +for node in teapot_latt.getNodes(): + (node_pos_start, node_pos_stop) = teapot_latt.getNodePositionsDict()[node] + if node_pos_start > pos_start and node_pos_stop < pos_stop: + print("node=", node.getName(), " type=", node.getType(), " pos=", node_pos_start, " L=", node.getLength()) + + +# ------------------------------ +# Main Bunch init +# ------------------------------ +b = Bunch() +print("Read Bunch.") +runName = "Benchmark_Foil" + +b.mass(0.93827231) +b.macroSize(1.0e1) +energy = 1.0 # Gev +# get initial bunch from ORBIT_MPI input +bunch_orbit_to_pyorbit(teapot_latt.getLength(), energy, os.path.join(script_dir, "Bm_KV_Uniform_10000"), b) +# b.readBunch("parts_in.dat") + +b.getSyncParticle().kinEnergy(energy) + +# =====track bunch through Collimator Node============ +paramsDict = {} +lostbunch = Bunch() +paramsDict["lostbunch"] = lostbunch +paramsDict["bunch"] = b + +lostbunch.addPartAttr("LostParticleAttributes") + +scatterchoice = 1 +foil.setScatterChoice(scatterchoice) +# collimator.trackBunch(b,paramsDict) +foil.track(paramsDict) + +# dump ORBIT_MPI bunch to compare results +bunch_pyorbit_to_orbit(teapot_latt.getLength(), b, "mainbunch.dat") +bunch_pyorbit_to_orbit(teapot_latt.getLength(), lostbunch, "lostbunch.dat") +lost_bunch = read_values_from_file("lostbunch.dat") +lostbunch.dumpBunch() +print("Stop.") + + +def test_getnHits_counts_and_accumulates(): + foil_node = TeapotFoilNode(-0.1, 0.1, -0.1, 0.1, 400) + foil_node.setScatterChoice(1) # Simple scattering + + bunch = Bunch() + bunch.mass(0.93827231) + bunch.getSyncParticle().kinEnergy(1.0) + + # Coordinates: x, xp, y, yp, z, dE + bunch.addParticle(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) # Inside + bunch.addParticle(0.2, 0.0, 0.0, 0.0, 0.0, 0.0) # Outside x + bunch.addParticle(0.0, 0.0, 0.2, 0.0, 0.0, 0.0) # Outside y + bunch.addParticle(.01, 0.0, .02, 0.0, 0.0, 0.0) + + params = {"bunch": bunch, "lostbunch": Bunch()} + + assert foil_node.getnHits() == 0 + print(f"Before tracking:{foil_node.getnHits()}") + foil_node.track(params) + print(f"Foil Hit after tracking:{foil_node.getnHits()}") #add particle coordinates after tracking through foil + assert foil_node.getnHits() == 2 + assert foil_node.getnHits() == 2 # Reading doesn't reset it + + foil_node.track(params) + print(f"After tracking a second time:{foil_node.getnHits()}") + assert foil_node.getnHits() == 4 + +os.remove("mainbunch.dat") +os.remove("lostbunch.dat")