diff --git a/examples/ESS_Linac/linac_solenoid_test/linac_solenoid_test.py b/examples/ESS_Linac/linac_solenoid_test/linac_solenoid_test.py new file mode 100755 index 00000000..efda2543 --- /dev/null +++ b/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.") + + diff --git a/examples/ESS_Linac/linac_solenoid_test/test_solenoid_lattice.xml b/examples/ESS_Linac/linac_solenoid_test/test_solenoid_lattice.xml new file mode 100644 index 00000000..7096a46f --- /dev/null +++ b/examples/ESS_Linac/linac_solenoid_test/test_solenoid_lattice.xml @@ -0,0 +1,20 @@ + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/py/orbit/py_linac/lattice/LinacAccNodes.py b/py/orbit/py_linac/lattice/LinacAccNodes.py index a807439e..8342d799 100755 --- a/py/orbit/py_linac/lattice/LinacAccNodes.py +++ b/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): """ diff --git a/py/orbit/py_linac/lattice/__init__.py b/py/orbit/py_linac/lattice/__init__.py index 5b61156b..cab3b3a7 100644 --- a/py/orbit/py_linac/lattice/__init__.py +++ b/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") diff --git a/py/orbit/py_linac/lattice_modifications/errors_modifications_lib.py b/py/orbit/py_linac/lattice_modifications/errors_modifications_lib.py index cbce8aa8..d6622b9f 100644 --- a/py/orbit/py_linac/lattice_modifications/errors_modifications_lib.py +++ b/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): """ diff --git a/py/orbit/py_linac/linac_parsers/sns_linac_lattice_factory.py b/py/orbit/py_linac/linac_parsers/sns_linac_lattice_factory.py index 5d85a293..d311750c 100644 --- a/py/orbit/py_linac/linac_parsers/sns_linac_lattice_factory.py +++ b/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")) diff --git a/src/spacecharge/SpaceChargeCalcUnifEllipse.cc b/src/spacecharge/SpaceChargeCalcUnifEllipse.cc index 5000129e..250ee81b 100644 --- a/src/spacecharge/SpaceChargeCalcUnifEllipse.cc +++ b/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;