Added getnHits to the Foil components of PyORBIT

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()

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;
+}

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:
 

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}
   	};
 

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")