Merge pull request #33 from shishlo/main

Author: azukov

examples/ESS_Linac/linac_solenoid_test/linac_solenoid_test.py

@@ -0,0 +1,123 @@
+#! /usr/bin/env python
+
+"""
+This script is the linac solenoid test.
+We will compare results with TEAPOT solenoid and will test the lattice parser. 
+"""
+
+import sys
+import math
+import random
+import time
+
+from orbit.py_linac.linac_parsers import SNS_LinacLatticeFactory
+
+from orbit.core.bunch import Bunch
+
+from orbit.lattice import AccLattice, AccNode, AccActionsContainer
+
+from orbit.py_linac.lattice import Quad,Solenoid
+
+from orbit.teapot import SolenoidTEAPOT
+
+#------------------------------------------------
+#    Start of Script
+#------------------------------------------------
+
+def printBunch(text,bunch):
+	"""
+	Print particles coordinates in the bunch.
+	"""
+	print (text)
+	b = bunch
+	for ind in range(bunch.getSize()):
+		(x, xp, y, yp, z, dE) = (b.x(ind),b.xp(ind),b.y(ind),b.yp(ind),b.z(ind),b.dE(ind))
+		st  = " %2d "%ind 
+		st += " %12.5g , %12.5g "%(x*1000., xp*1000.)
+		st += " %12.5g , %12.5g "%(y*1000., yp*1000.)
+		st += " %12.5g , %12.5g "%(z, dE*1000.)
+		print (st)
+	print ("==================================")
+
+names = ["TEST_Seq",]
+
+sns_linac_factory = SNS_LinacLatticeFactory()
+sns_linac_factory.setMaxDriftLength(0.5)
+
+# ---- the XML file name with the structure
+xml_file_name = "test_solenoid_lattice.xml"
+
+# ---- make lattice from XML file
+accLattice = sns_linac_factory.getLinacAccLattice(names, xml_file_name)
+
+print ("============ Test Lattice =============")
+for ind,node in enumerate(accLattice.getNodes()):
+	st  = "ind= %2d "%ind
+	st += " node= %35s "%node.getName()
+	st += " length[m] = %6.4f "%node.getLength()
+	st += " nParts= %2d "%node.getnParts()
+	st += " type= %10s "%node.getType()
+	print (st)
+print ("========================================")
+
+
+node_soln = accLattice.getNodesOfClass(Solenoid)[0]
+node_soln_ind = accLattice.getNodeIndex(node_soln)
+print ("Solenoid field B=",node_soln.getParam("B"))
+
+#---- Results should be the same for any parts in the solenoid node
+nSolParts = 10
+node_soln.setnParts(nSolParts)
+print (" nParts= %2d "%node_soln.getnParts())
+
+bunch_init = Bunch()
+syncPart = bunch_init.getSyncParticle()
+bunch_init.mass(0.9382723)
+bunch_init.charge(+1.0)
+syncPart.kinEnergy(0.003)
+
+(x, xp, y, yp, z, dE) = (0.,0.,0.,0.,0.,0.)
+bunch_init.addParticle(x, xp, y, yp, z, dE)
+
+(x, xp, y, yp, z, dE) = (0.001,0.001,-0.001,-0.001,0.1,0.00002)
+bunch_init.addParticle(x, xp, y, yp, z, dE)
+bunch_init.addParticle(x/2, xp/2, y/2, yp/2, z/2, dE/2)
+
+(x, xp, y, yp, z, dE) = (-0.001,-0.001,+0.001,+0.001,-0.1,-0.00002)
+bunch_init.addParticle(x, xp, y, yp, z, dE)
+bunch_init.addParticle(x/2, xp/2, y/2, yp/2, z/2, dE/2)
+
+#---- set up design
+accLattice.trackDesignBunch(bunch_init)
+
+#---- track bunch to the solenoid start
+accLattice.trackBunch(bunch_init,None,None,0,node_soln_ind-1)
+
+printBunch("========= Initial Bunch: before Solenoid ",bunch_init)
+
+bunch = Bunch()
+bunch_init.copyBunchTo(bunch)
+
+#---- track bunch through the solenoid
+accLattice.trackBunch(bunch,None,None,node_soln_ind,node_soln_ind)
+
+printBunch("========= Bunch After: after Solenoid ",bunch)
+
+#----------------------------------------------
+# Now let's create TEAPOT solenoid and track
+# the same bunch through it.
+#----------------------------------------------
+
+soln_teapot = SolenoidTEAPOT("solenoid_teapot")
+soln_teapot.setParam("B",node_soln.getParam("B"))
+soln_teapot.setLength(node_soln.getLength())
+
+bunch = Bunch()
+bunch_init.copyBunchTo(bunch)
+soln_teapot.trackBunch(bunch)
+
+printBunch("========= Bunch After: after TEAPOT Solenoid - should be the same",bunch)
+
+print ("Stop.")
+
+

examples/ESS_Linac/linac_solenoid_test/test_solenoid_lattice.xml

@@ -0,0 +1,20 @@
+<?xml version="1.0" ?>
+<ess>
+	<TEST_Seq bpmFrequency="7.0442E8" length="3" name="MEBT">
+		<accElement length="0.0" name="TEST_START" pos="0.0" type="MARKER">
+			<parameters/>
+		</accElement>
+		<accElement length="0.10" name="QV01" pos="0.5" type="QUAD">
+			<parameters aperture="0.5" aprt_type="1" field="-10.0"/>
+		</accElement>
+		<accElement length="1" name="SL00" pos="1.5" type="SOLENOID">
+			<parameters aperture="0.5" aprt_type="1" B="10.0"/>
+		</accElement>
+		<accElement length="0.10" name="QV01" pos="2.5" type="QUAD">
+			<parameters aperture="0.5" aprt_type="1" field="+10.0"/>
+		</accElement>
+		<accElement length="0.0" name="TEST_EDN" pos="2.9999999" type="MARKER">
+			<parameters/>
+		</accElement>
+	</TEST_Seq>
+</ess>

py/orbit/py_linac/lattice/LinacAccNodes.py

@@ -849,6 +849,32 @@ def track(self, paramsDict):
         self.tracking_module.kick(bunch, kickX, kickY, 0.0)
         self.tracking_module.drift(bunch, length / 2.0)
 
+class Solenoid(BaseLinacNode):
+    """
+    Solenoid TEAPOT based element.
+    """
+
+    def __init__(self, name="solenoid no name"):
+        """
+        Constructor. Creates the Solenoid TEAPOT based element.
+        """
+        BaseLinacNode.__init__(self, name)
+        self.setType("solenoid")
+        self.addParam("B", 0.0)
+
+    def track(self, paramsDict):
+        """
+        The Solenoid TEAPOT class implementation of the
+        AccNodeBunchTracker class track(probe) method.
+        """
+        index = self.getActivePartIndex()
+        length = self.getLength(index)
+        B = self.getParam("B")
+        bunch = paramsDict["bunch"]
+        useCharge = 1
+        if "useCharge" in paramsDict:
+            useCharge = paramsDict["useCharge"]
+        TPB.soln(bunch, length, B, useCharge)
 
 class AbstractRF_Gap(BaseLinacNode):
     """

py/orbit/py_linac/lattice/__init__.py

@@ -10,6 +10,7 @@
 
 from orbit.py_linac.lattice.LinacAccNodes import BaseLinacNode, LinacNode, LinacMagnetNode
 from orbit.py_linac.lattice.LinacAccNodes import MarkerLinacNode, Drift, Quad, AbstractRF_Gap, Bend
+from orbit.py_linac.lattice.LinacAccNodes import Solenoid
 from orbit.py_linac.lattice.LinacAccNodes import DCorrectorH, DCorrectorV
 from orbit.py_linac.lattice.LinacAccNodes import ThickKick
 
@@ -53,6 +54,7 @@
 __all__.append("DCorrectorV")
 __all__.append("ThickKick")
 __all__.append("Bend")
+__all__.append("Solenoid")
 
 
 __all__.append("LinacApertureNode")

py/orbit/py_linac/lattice_modifications/errors_modifications_lib.py

@@ -485,7 +485,8 @@ def restoreFields(self):
 
     def setGaussDistributedRealtiveErrors(self, relative_error, cut_off_level=3.0, comm=mpi_comm.MPI_COMM_WORLD):
         """
-        Sets the random generated error field for all quads.
+        Sets the random generated error fields for all quads.
+        Changes will be different for all quads.
         """
         for [quad, field_init] in self.quad_and_field_arr:
             rel_err = random.gauss(0.0, relative_error)
@@ -496,6 +497,19 @@ def setGaussDistributedRealtiveErrors(self, relative_error, cut_off_level=3.0, c
             field = field_init * (1.0 + rel_err)
             quad.setParam("dB/dr", field)
 
+    def setGaussDistributedRealtiveErrorToGroup(self,relative_error,cut_off_level = 3.0,comm = mpi_comm.MPI_COMM_WORLD):
+        """
+        Sets the random generated error fields for all quads in the string.
+        Changes will be the same for all quads. We assume one PS for all quads.
+        """
+        rel_err = random.gauss(0.,relative_error)
+        while(abs(rel_err) > abs(relative_error)*cut_off_level):
+            rel_err = random.gauss(0.,relative_error)
+        main_rank = 0
+        rel_err = orbit_mpi.MPI_Bcast(rel_err,mpi_datatype.MPI_DOUBLE,main_rank,comm)
+        for [quad,field_init] in self.quad_and_field_arr:
+            field = field_init*(1.0 + rel_err)
+            quad.setParam("dB/dr",field)
 
 class BendFieldNodesModification(ErrorForNodesModification):
     """

py/orbit/py_linac/linac_parsers/sns_linac_lattice_factory.py

@@ -19,6 +19,7 @@
 
 from orbit.py_linac.lattice import BaseLinacNode, LinacNode, LinacMagnetNode, MarkerLinacNode, Drift, Quad, AbstractRF_Gap, Bend
 from orbit.py_linac.lattice import DCorrectorH, DCorrectorV, ThickKick
+from orbit.py_linac.lattice import Solenoid
 from orbit.py_linac.lattice import RF_Cavity, Sequence
 from orbit.py_linac.lattice import BaseRF_Gap
 
@@ -201,6 +202,12 @@ def positionComp(node_da):
                         accNode.setnParts(2 * int(accNode.getLength() / self.maxDriftLength + 1.5 - 1.0e-12))
                     accNode.setParam("pos", node_pos)
                     accSeq.addNode(accNode)
+                elif node_type == "SOLENOID":
+                    accNode = Solenoid(node_da.stringValue("name"))
+                    accNode.setParam("B",params_da.doubleValue("B"))
+                    accNode.setLength(node_length)
+                    accNode.setParam("pos", node_pos)
+                    accSeq.addNode(accNode)
                 # ------------RF_Gap-----------------
                 elif node_type == "RFGAP":
                     accNode = BaseRF_Gap(node_da.stringValue("name"))

src/spacecharge/SpaceChargeCalcUnifEllipse.cc

@@ -146,7 +146,12 @@ void SpaceChargeCalcUnifEllipse::bunchAnalysis(Bunch* bunch){
 	ORBIT_MPI_Allreduce(coord_avg,coord_avg_out,7,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm);
 
 	total_macrosize = coord_avg_out[6];
-	if(total_macrosize == 0.) return;
+	if(total_macrosize == 0.){
+	  //free resources
+	  OrbitUtils::BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index0);
+	  OrbitUtils::BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index1);		
+		return;
+	}
 
 	//calculate the parameters of the biggest ellipse
 	x_center = coord_avg_out[0]/total_macrosize;