From 464491947e4f582295191a9cd61a7b411a6c1952 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Fri, 4 Sep 2026 23:52:44 -0400 Subject: [PATCH 1/6] First version of 2D uniform ellipsoid space charge calculator --- py/orbit/space_charge/sc2p5d/__init__.py | 4 + py/orbit/space_charge/sc2p5d/scAccNodes.py | 12 + .../sc2p5d/scLatticeModifications.py | 10 +- py/orbit/space_charge/sc3d/__init__.py | 8 +- py/orbit/space_charge/sc3d/scAccNodes.py | 2 +- .../sc3d/scLatticeModifications.py | 10 +- src/meson.build | 2 +- src/spacecharge/SpaceChargeCalcUnifEllipse.cc | 747 +++++++++++------- src/spacecharge/SpaceChargeCalcUnifEllipse.hh | 136 ++-- .../UniformEllipsoidFieldCalculator.cc | 508 ++++++------ .../UniformEllipsoidFieldCalculator.hh | 149 ++-- src/spacecharge/wrap_spacecharge.cc | 4 +- .../wrap_spacechargecalc_unif_ellipse.cc | 247 ++++++ ...h => wrap_spacechargecalc_unif_ellipse.hh} | 6 +- .../wrap_spacechargecalc_uniform_ellipse.cc | 197 ----- ...wrap_uniform_ellipsoid_field_calculator.cc | 257 +++--- tests/py/orbit/test_sc_unif_ellipse.py | 77 ++ tests/py/orbit/test_sc_unif_ellipse_nodes.py | 32 + 18 files changed, 1431 insertions(+), 977 deletions(-) create mode 100644 src/spacecharge/wrap_spacechargecalc_unif_ellipse.cc rename src/spacecharge/{wrap_spacechargecalc_uniform_ellipse.hh => wrap_spacechargecalc_unif_ellipse.hh} (72%) delete mode 100644 src/spacecharge/wrap_spacechargecalc_uniform_ellipse.cc create mode 100644 tests/py/orbit/test_sc_unif_ellipse.py create mode 100644 tests/py/orbit/test_sc_unif_ellipse_nodes.py diff --git a/py/orbit/space_charge/sc2p5d/__init__.py b/py/orbit/space_charge/sc2p5d/__init__.py index 4ea42d7d..15ad9b53 100644 --- a/py/orbit/space_charge/sc2p5d/__init__.py +++ b/py/orbit/space_charge/sc2p5d/__init__.py @@ -6,11 +6,15 @@ from orbit.space_charge.sc2p5d.scAccNodes import SC2p5Drb_AccNode from orbit.space_charge.sc2p5d.scAccNodes import SC2p5D_AccNode +from orbit.space_charge.sc2p5d.scAccNodes import SCUnifEllipse2D_AccNode from orbit.space_charge.sc2p5d.scLatticeModifications import setSC2p5DAccNodes from orbit.space_charge.sc2p5d.scLatticeModifications import setSC2p5DrbAccNodes +from orbit.space_charge.sc2p5d.scLatticeModifications import setSCUnifEllipse2DAccNodes __all__ = [] __all__.append("SC2p5Drb_AccNode") __all__.append("SC2p5D_AccNode") +__all__.append("SCUnifEllipse2D_AccNode") __all__.append("setSC2p5DAccNodes") __all__.append("setSC2p5DrbAccNodes") +__all__.append("setSCUnifEllipse2DAccNodes") diff --git a/py/orbit/space_charge/sc2p5d/scAccNodes.py b/py/orbit/space_charge/sc2p5d/scAccNodes.py index bf42ac90..9ea35925 100644 --- a/py/orbit/space_charge/sc2p5d/scAccNodes.py +++ b/py/orbit/space_charge/sc2p5d/scAccNodes.py @@ -88,3 +88,15 @@ def track(self, paramsDict): return bunch = paramsDict["bunch"] self.sc_calculator.trackBunch(bunch, self.sc_length, self.pipe_radius) + + +class SCUnifEllipse2D_AccNode(SC_Base_AccNode): + def __init__(self, sc_calculator, name="no name"): + SC_Base_AccNode.__init__(self, sc_calculator, name) + self.setType("UnifEllsSC2D") + + def track(self, paramsDict): + if self.switcher != True: + return + bunch = paramsDict["bunch"] + self.sc_calculator.trackBunch(bunch, self.sc_length) diff --git a/py/orbit/space_charge/sc2p5d/scLatticeModifications.py b/py/orbit/space_charge/sc2p5d/scLatticeModifications.py index dad47309..e4d5d81d 100644 --- a/py/orbit/space_charge/sc2p5d/scLatticeModifications.py +++ b/py/orbit/space_charge/sc2p5d/scLatticeModifications.py @@ -3,7 +3,7 @@ """ # import SC acc. nodes -from orbit.space_charge.sc2p5d import SC2p5D_AccNode, SC2p5Drb_AccNode +from orbit.space_charge.sc2p5d import SC2p5D_AccNode, SC2p5Drb_AccNode, SCUnifEllipse2D_AccNode # import general accelerator elements and lattice from orbit.lattice import AccLattice, AccNode, AccActionsContainer, AccNodeBunchTracker @@ -45,3 +45,11 @@ def setSC2p5DrbAccNodes(lattice, sc_path_length_min, space_charge_calculator, pi # initialize the lattice lattice.initialize() return scNodes_arr + + +def setSCUnifEllipse2DAccNodes(lattice, sc_path_length_min, space_charge_calculator): + scNodes_arr = setSC_General_AccNodes(lattice, sc_path_length_min, space_charge_calculator, SCUnifEllipse2D_AccNode) + for scNode in scNodes_arr: + scNode.setName(scNode.getName() + "SCUnifEllipse2D") + lattice.initialize() + return scNodes_arr diff --git a/py/orbit/space_charge/sc3d/__init__.py b/py/orbit/space_charge/sc3d/__init__.py index 96003d2b..f00d0ce1 100644 --- a/py/orbit/space_charge/sc3d/__init__.py +++ b/py/orbit/space_charge/sc3d/__init__.py @@ -5,12 +5,12 @@ ## from orbit.space_charge.sc3d.scAccNodes import SC3D_AccNode -from orbit.space_charge.sc3d.scAccNodes import SC_UniformEllipses_AccNode +from orbit.space_charge.sc3d.scAccNodes import SCUnifEllipse_AccNode from orbit.space_charge.sc3d.scLatticeModifications import setSC3DAccNodes -from orbit.space_charge.sc3d.scLatticeModifications import setUniformEllipsesSCAccNodes +from orbit.space_charge.sc3d.scLatticeModifications import setSCUnifEllipseAccNodes __all__ = [] __all__.append("SC3D_AccNode") -__all__.append("SC_UniformEllipses_AccNode") +__all__.append("SCUnifEllipse_AccNode") __all__.append("setSC3DAccNodes") -__all__.append("setUniformEllipsesSCAccNodes") +__all__.append("setSCUnifEllipseAccNodes") diff --git a/py/orbit/space_charge/sc3d/scAccNodes.py b/py/orbit/space_charge/sc3d/scAccNodes.py index 3051883d..344eb6b2 100644 --- a/py/orbit/space_charge/sc3d/scAccNodes.py +++ b/py/orbit/space_charge/sc3d/scAccNodes.py @@ -38,7 +38,7 @@ def track(self, paramsDict): self.sc_calculator.trackBunch(bunch, self.sc_length) -class SC_UniformEllipses_AccNode(SC_Base_AccNode): +class SCUnifEllipse_AccNode(SC_Base_AccNode): """ The subclass of the AccNodeBunchTracker class. It uses SpaceChargeCalcUnifEllipse wrapper for the c++ space charge calculator. """ diff --git a/py/orbit/space_charge/sc3d/scLatticeModifications.py b/py/orbit/space_charge/sc3d/scLatticeModifications.py index 800a4e33..c0b6a25f 100644 --- a/py/orbit/space_charge/sc3d/scLatticeModifications.py +++ b/py/orbit/space_charge/sc3d/scLatticeModifications.py @@ -3,7 +3,7 @@ """ # import SC acc. nodes -from orbit.space_charge.sc3d import SC3D_AccNode, SC_UniformEllipses_AccNode +from orbit.space_charge.sc3d import SC3D_AccNode, SCUnifEllipse_AccNode # import general accelerator elements and lattice from orbit.lattice import AccLattice, AccNode, AccActionsContainer, AccNodeBunchTracker @@ -27,16 +27,16 @@ def setSC3DAccNodes(lattice, sc_path_length_min, space_charge_calculator): return scNodes_arr -def setUniformEllipsesSCAccNodes(lattice, sc_path_length_min, space_charge_calculator): +def setSCUnifEllipseAccNodes(lattice, sc_path_length_min, space_charge_calculator): """ - It will put a set of a space charge SC_UniformEllipses_AccNode into the lattice as child nodes of the first level accelerator nodes. + It will put a set of SCUnifEllipse_AccNode nodes into the lattice as child nodes of the first level accelerator nodes. The SC nodes will be inserted at the beginning of a particular part of the first level AccNode element. The distance between SC nodes should be more than sc_path_length_min. The function will return the array of SC nodes as a convenience for the user. """ - scNodes_arr = setSC_General_AccNodes(lattice, sc_path_length_min, space_charge_calculator, SC_UniformEllipses_AccNode) + scNodes_arr = setSC_General_AccNodes(lattice, sc_path_length_min, space_charge_calculator, SCUnifEllipse_AccNode) for scNode in scNodes_arr: - scNode.setName(scNode.getName() + "UnifEllsSC") + scNode.setName(scNode.getName() + "SCUnifEllipse") # initialize the lattice lattice.initialize() return scNodes_arr diff --git a/src/meson.build b/src/meson.build index c223553d..e036bd7d 100644 --- a/src/meson.build +++ b/src/meson.build @@ -187,7 +187,7 @@ sources = files([ 'spacecharge/wrap_grid3D.cc', 'spacecharge/BaseBoundary2D.cc', 'spacecharge/wrap_grid2D.cc', - 'spacecharge/wrap_spacechargecalc_uniform_ellipse.cc', + 'spacecharge/wrap_spacechargecalc_unif_ellipse.cc', 'spacecharge/wrap_spacechargecalc3d.cc', 'spacecharge/UniformEllipsoidFieldCalculator.cc', 'spacecharge/wrap_spacechargeforcecalc2p5d.cc', diff --git a/src/spacecharge/SpaceChargeCalcUnifEllipse.cc b/src/spacecharge/SpaceChargeCalcUnifEllipse.cc index 250ee81b..d0e7384a 100644 --- a/src/spacecharge/SpaceChargeCalcUnifEllipse.cc +++ b/src/spacecharge/SpaceChargeCalcUnifEllipse.cc @@ -20,293 +20,502 @@ #include "ParticleMacroSize.hh" -#include -#include #include +#include +#include using namespace OrbitUtils; -SpaceChargeCalcUnifEllipse::SpaceChargeCalcUnifEllipse(int nEllipses_in): CppPyWrapper(NULL) -{ - nEllipses = nEllipses_in; - ellipsoidCalc_arr = new UniformEllipsoidFieldCalculator*[nEllipses]; - for(int ie = 0; ie < nEllipses; ie++){ - ellipsoidCalc_arr[ie] = new UniformEllipsoidFieldCalculator(); - } - macroSizesEll_arr = (double* ) malloc (sizeof(double)*nEllipses); - macroSizesEll_MPI_arr = (double* ) malloc (sizeof(double)*nEllipses); - for(int ie = 0; ie < nEllipses; ie++){ - macroSizesEll_arr[ie] = 0.; - macroSizesEll_MPI_arr[ie] = 0.; - } +SpaceChargeCalcUnifEllipse::SpaceChargeCalcUnifEllipse(int nEllipses_in) : CppPyWrapper(NULL) { + nEllipses = nEllipses_in; + ellipsoidCalc_arr = new UniformEllipsoidFieldCalculator *[nEllipses]; + for (int ie = 0; ie < nEllipses; ie++) { + ellipsoidCalc_arr[ie] = new UniformEllipsoidFieldCalculator(); + } + macroSizesEll_arr = (double *)malloc(sizeof(double) * nEllipses); + macroSizesEll_MPI_arr = (double *)malloc(sizeof(double) * nEllipses); + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_arr[ie] = 0.; + macroSizesEll_MPI_arr[ie] = 0.; + } } - -SpaceChargeCalcUnifEllipse::~SpaceChargeCalcUnifEllipse(){ - for(int ie = 0; ie < nEllipses; ie++){ - if(ellipsoidCalc_arr[ie]->getPyWrapper() != NULL){ - Py_DECREF(ellipsoidCalc_arr[ie]->getPyWrapper()); - } else { - delete ellipsoidCalc_arr[ie]; - } - } - delete [] ellipsoidCalc_arr; - - free(macroSizesEll_arr); - free(macroSizesEll_MPI_arr); +SpaceChargeCalcUnifEllipse::~SpaceChargeCalcUnifEllipse() { + for (int ie = 0; ie < nEllipses; ie++) { + if (ellipsoidCalc_arr[ie]->getPyWrapper() != NULL) { + Py_DECREF(ellipsoidCalc_arr[ie]->getPyWrapper()); + } else { + delete ellipsoidCalc_arr[ie]; + } + } + delete[] ellipsoidCalc_arr; + + free(macroSizesEll_arr); + free(macroSizesEll_MPI_arr); } - -void SpaceChargeCalcUnifEllipse::trackBunch(Bunch* bunch, double length){ - - int nPartsGlobal = bunch->getSizeGlobal(); - if(nPartsGlobal < 3) return; - - SyncPart* syncPart = bunch->getSyncPart(); - double beta = syncPart->getBeta(); - double gamma = syncPart->getGamma(); - - for(int ie = 0; ie < nEllipses; ie++){ - ellipsoidCalc_arr[ie]->setQ(0.); - } - - //analyse the bunch and make the ellipsoid filed sources - this->bunchAnalysis(bunch); - - //if there is nothing we give up - if(total_macrosize == 0.) return; - - double trans_factor = length*bunch->getClassicalRadius()/(pow(beta,2)*pow(gamma,2)); - double long_factor = length*bunch->getClassicalRadius()*bunch->getMass(); - - double x,y,z,ex,ey,ez; - for (int i = 0, n = bunch->getSize(); i < n; i++){ - x = bunch->x(i) - x_center; - y = bunch->y(i) - y_center; - z = (bunch->z(i) - z_center)*gamma; - this->calculateField(x,y,z,ex,ey,ez); - //calculate momentum kicks - bunch->xp(i) += ex * trans_factor; - bunch->yp(i) += ey * trans_factor; - bunch->dE(i) += ez * long_factor; - } +void SpaceChargeCalcUnifEllipse::trackBunch(Bunch *bunch, double length) { + + int nPartsGlobal = bunch->getSizeGlobal(); + if (nPartsGlobal < 3) + return; + + SyncPart *syncPart = bunch->getSyncPart(); + double beta = syncPart->getBeta(); + double gamma = syncPart->getGamma(); + + for (int ie = 0; ie < nEllipses; ie++) { + ellipsoidCalc_arr[ie]->setQ(0.); + } + + // analyse the bunch and make the ellipsoid filed sources + this->bunchAnalysis(bunch); + + // if there is nothing we give up + if (total_macrosize == 0.) + return; + + double trans_factor = length * bunch->getClassicalRadius() / (pow(beta, 2) * pow(gamma, 2)); + double long_factor = length * bunch->getClassicalRadius() * bunch->getMass(); + + double x, y, z, ex, ey, ez; + for (int i = 0, n = bunch->getSize(); i < n; i++) { + x = bunch->x(i) - x_center; + y = bunch->y(i) - y_center; + z = (bunch->z(i) - z_center) * gamma; + this->calculateField(x, y, z, ex, ey, ez); + // calculate momentum kicks + bunch->xp(i) += ex * trans_factor; + bunch->yp(i) += ey * trans_factor; + bunch->dE(i) += ez * long_factor; + } } /** Analyses the bunch and sets up the ellipsoid filed sources */ -void SpaceChargeCalcUnifEllipse::bunchAnalysis(Bunch* bunch){ - - //average values for x,y,z,x2,y2,z2 and total macrosize - int buff_index0 = 0; - int buff_index1 = 0; - double* coord_avg = BufferStore::getBufferStore()->getFreeDoubleArr(buff_index0,7); - double* coord_avg_out = BufferStore::getBufferStore()->getFreeDoubleArr(buff_index1,7); - for (int i = 0; i < 7; i++){ - coord_avg[i] = 0.; - } - - //caluclate limits and averages - double** partArr=bunch->coordArr(); - double* coordArr = NULL; - bunch->compress(); - double** part_coord_arr = bunch->coordArr(); - int has_msize = bunch->hasParticleAttributes("macrosize"); - if(has_msize > 0){ - ParticleMacroSize* macroSizeAttr = (ParticleMacroSize*) bunch->getParticleAttributes("macrosize"); - double m_size = 0.; - for(int ip = 0, n = bunch->getSize(); ip < n; ip++){ - m_size = macroSizeAttr->macrosize(ip); - coordArr = partArr[ip]; - coord_avg[0] += m_size*coordArr[0]; - coord_avg[1] += m_size*coordArr[2]; - coord_avg[2] += m_size*coordArr[4]; - coord_avg[3] += m_size*coordArr[0]*coordArr[0]; - coord_avg[4] += m_size*coordArr[2]*coordArr[2]; - coord_avg[5] += m_size*coordArr[4]*coordArr[4]; - coord_avg[6] += m_size; - } - } else { - double m_size = bunch->getMacroSize(); - int nParts = bunch->getSize(); - coord_avg[6] = m_size*nParts; - for(int ip = 0; ip < nParts; ip++){ - coordArr = partArr[ip]; - coord_avg[0] += coordArr[0]; - coord_avg[1] += coordArr[2]; - coord_avg[2] += coordArr[4]; - coord_avg[3] += coordArr[0]*coordArr[0]; - coord_avg[4] += coordArr[2]*coordArr[2]; - coord_avg[5] += coordArr[4]*coordArr[4]; - } - for (int i = 0; i < 6; i++){ - coord_avg[i] *= m_size; - } - } - - //calculates sum over all CPUs - 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.){ - //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; - y_center = coord_avg_out[1]/total_macrosize; - z_center = coord_avg_out[2]/total_macrosize; - x2_avg = fabs(coord_avg_out[3]/total_macrosize - x_center*x_center); - y2_avg = fabs(coord_avg_out[4]/total_macrosize - y_center*y_center); - z2_avg = fabs(coord_avg_out[5]/total_macrosize - z_center*z_center); - a2_ellips = 5.0*x2_avg; - b2_ellips = 5.0*y2_avg; - c2_ellips = 5.0*z2_avg; - a_ellips = sqrt(a2_ellips); - b_ellips = sqrt(b2_ellips); - c_ellips = sqrt(c2_ellips); - - //std::cout<<"debug a_ellips="<< a_ellips <<" b_ellips="<< b_ellips <<" c_ellips="<< c_ellips <setUnusedDoubleArr(buff_index0); - OrbitUtils::BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index1); - - //check if the beam size is not zero - if( x2_avg == 0. || y2_avg == 0.|| z2_avg == 0.){ - int rank = 0; - ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); - if(rank == 0){ - std::cerr << "SpaceChargeCalcUnifEllipse::bunchAnalysis(bunch,...)" << std::endl - << "The bunch coords min and max sizes are wrong! Cannot calculate space charge!" << std::endl - <<" x2_rms="<< x2_avg << std::endl - <<" y2_rms="<< y2_avg << std::endl - <<" z2_rms="<< z2_avg << std::endl - << "Stop."<< std::endl; - } - ORBIT_MPI_Finalize(); - } - - //relativistic factor gamma - double gamma = bunch->getSyncPart()->getGamma(); - - //if we have only one ellipse we should not distribute anything - if(nEllipses == 1){ - macroSizesEll_arr[0] = total_macrosize; - double r_max = a_ellips; - if(r_max < b_ellips) r_max = b_ellips; - if(r_max < c_ellips*gamma) r_max = c_ellips*gamma; - ellipsoidCalc_arr[0]->setEllipsoid(a_ellips,b_ellips,c_ellips*gamma,10.*r_max); - ellipsoidCalc_arr[0]->setQ(macroSizesEll_arr[0]); - return; - } - - //find the distribution of the macrosizes between nEllipses - for(int ie = 0; ie < nEllipses; ie++){ - macroSizesEll_arr[ie] = 0.; - } - - double pos = 0.; - int pos_index = 0; - if(has_msize > 0){ - ParticleMacroSize* macroSizeAttr = (ParticleMacroSize*) bunch->getParticleAttributes("macrosize"); - double m_size = 0.; - for(int ip = 0, n = bunch->getSize(); ip < n; ip++){ - m_size = macroSizeAttr->macrosize(ip); - coordArr = partArr[ip]; - pos = sqrt(coordArr[0]*coordArr[0]/a2_ellips + coordArr[2]*coordArr[2]/b2_ellips + coordArr[4]*coordArr[4]/c2_ellips); - pos_index = int(pos*nEllipses); - if(pos_index < 0) pos_index = 0; - if(pos_index >= nEllipses) pos_index = nEllipses - 1; - macroSizesEll_arr[pos_index] += m_size; - } - } else { - double m_size = bunch->getMacroSize(); - int nParts = bunch->getSize(); - for(int ip = 0, n = bunch->getSize(); ip < n; ip++){ - coordArr = partArr[ip]; - pos = sqrt(coordArr[0]*coordArr[0]/a2_ellips + coordArr[2]*coordArr[2]/b2_ellips + coordArr[4]*coordArr[4]/c2_ellips); - pos_index = int(pos*nEllipses) - 1; - if(pos_index < 0) pos_index = 0; - if(pos_index >= nEllipses) pos_index = nEllipses - 1; - macroSizesEll_arr[pos_index] += m_size; - } - } - //calculates sum over all CPUs - ORBIT_MPI_Allreduce(macroSizesEll_arr,macroSizesEll_MPI_arr,nEllipses,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - for(int ie = 0; ie < nEllipses; ie++){ - macroSizesEll_arr[ie] = macroSizesEll_MPI_arr[ie]; - //std::cout<<"debug 0 ie ="<< ie <<" macrosize="<< macroSizesEll_MPI_arr[ie] << std::endl; - } - //calculate the relative volume density in each region. This density is a sum of all elipsoids - for(int ie = 0; ie < nEllipses; ie++){ - macroSizesEll_MPI_arr[ie] /= ((ie+2)*(ie+2)*(ie+2) - (ie+1)*(ie+1)*(ie+1)); - //std::cout<<"debug 1 ie ="<< ie <<" macrosize="<< macroSizesEll_MPI_arr[ie] << std::endl; - } - //calculate the density for each elipsoid - double rho_sum = 0.; - for(int ie = (nEllipses-1); ie >= 0; ie--){ - macroSizesEll_MPI_arr[ie] -= rho_sum; - rho_sum += macroSizesEll_MPI_arr[ie]; - //std::cout<<"debug 2 ie ="<< ie <<" macrosize="<< macroSizesEll_MPI_arr[ie] << " rho_sum="<< rho_sum < 0) { + ParticleMacroSize *macroSizeAttr = (ParticleMacroSize *)bunch->getParticleAttributes("macrosize"); + double m_size = 0.; + for (int ip = 0, n = bunch->getSize(); ip < n; ip++) { + m_size = macroSizeAttr->macrosize(ip); + coordArr = partArr[ip]; + coord_avg[0] += m_size * coordArr[0]; + coord_avg[1] += m_size * coordArr[2]; + coord_avg[2] += m_size * coordArr[4]; + coord_avg[3] += m_size * coordArr[0] * coordArr[0]; + coord_avg[4] += m_size * coordArr[2] * coordArr[2]; + coord_avg[5] += m_size * coordArr[4] * coordArr[4]; + coord_avg[6] += m_size; + } + } else { + double m_size = bunch->getMacroSize(); + int nParts = bunch->getSize(); + coord_avg[6] = m_size * nParts; + for (int ip = 0; ip < nParts; ip++) { + coordArr = partArr[ip]; + coord_avg[0] += coordArr[0]; + coord_avg[1] += coordArr[2]; + coord_avg[2] += coordArr[4]; + coord_avg[3] += coordArr[0] * coordArr[0]; + coord_avg[4] += coordArr[2] * coordArr[2]; + coord_avg[5] += coordArr[4] * coordArr[4]; + } + for (int i = 0; i < 6; i++) { + coord_avg[i] *= m_size; + } + } + + // calculates sum over all CPUs + 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.) { + // 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; + y_center = coord_avg_out[1] / total_macrosize; + z_center = coord_avg_out[2] / total_macrosize; + x2_avg = fabs(coord_avg_out[3] / total_macrosize - x_center * x_center); + y2_avg = fabs(coord_avg_out[4] / total_macrosize - y_center * y_center); + z2_avg = fabs(coord_avg_out[5] / total_macrosize - z_center * z_center); + a2_ellips = 5.0 * x2_avg; + b2_ellips = 5.0 * y2_avg; + c2_ellips = 5.0 * z2_avg; + a_ellips = sqrt(a2_ellips); + b_ellips = sqrt(b2_ellips); + c_ellips = sqrt(c2_ellips); + + // std::cout<<"debug a_ellips="<< a_ellips <<" b_ellips="<< b_ellips <<" c_ellips="<< c_ellips <setUnusedDoubleArr(buff_index0); + OrbitUtils::BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index1); + + // check if the beam size is not zero + if (x2_avg == 0. || y2_avg == 0. || z2_avg == 0.) { + int rank = 0; + ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) { + std::cerr << "SpaceChargeCalcUnifEllipse::bunchAnalysis(bunch,...)" << std::endl + << "The bunch coords min and max sizes are wrong! Cannot calculate space charge!" << std::endl + << " x2_rms=" << x2_avg << std::endl + << " y2_rms=" << y2_avg << std::endl + << " z2_rms=" << z2_avg << std::endl + << "Stop." << std::endl; + } + ORBIT_MPI_Finalize(); + } + + // relativistic factor gamma + double gamma = bunch->getSyncPart()->getGamma(); + + // if we have only one ellipse we should not distribute anything + if (nEllipses == 1) { + macroSizesEll_arr[0] = total_macrosize; + double r_max = a_ellips; + if (r_max < b_ellips) + r_max = b_ellips; + if (r_max < c_ellips * gamma) + r_max = c_ellips * gamma; + ellipsoidCalc_arr[0]->setEllipsoid(a_ellips, b_ellips, c_ellips * gamma, 10. * r_max); + ellipsoidCalc_arr[0]->setQ(macroSizesEll_arr[0]); + return; + } + + // find the distribution of the macrosizes between nEllipses + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_arr[ie] = 0.; + } + + double pos = 0.; + int pos_index = 0; + if (has_msize > 0) { + ParticleMacroSize *macroSizeAttr = (ParticleMacroSize *)bunch->getParticleAttributes("macrosize"); + double m_size = 0.; + for (int ip = 0, n = bunch->getSize(); ip < n; ip++) { + m_size = macroSizeAttr->macrosize(ip); + coordArr = partArr[ip]; + pos = sqrt(coordArr[0] * coordArr[0] / a2_ellips + coordArr[2] * coordArr[2] / b2_ellips + coordArr[4] * coordArr[4] / c2_ellips); + pos_index = int(pos * nEllipses); + if (pos_index < 0) + pos_index = 0; + if (pos_index >= nEllipses) + pos_index = nEllipses - 1; + macroSizesEll_arr[pos_index] += m_size; + } + } else { + double m_size = bunch->getMacroSize(); + int nParts = bunch->getSize(); + for (int ip = 0, n = bunch->getSize(); ip < n; ip++) { + coordArr = partArr[ip]; + pos = sqrt(coordArr[0] * coordArr[0] / a2_ellips + coordArr[2] * coordArr[2] / b2_ellips + coordArr[4] * coordArr[4] / c2_ellips); + pos_index = int(pos * nEllipses) - 1; + if (pos_index < 0) + pos_index = 0; + if (pos_index >= nEllipses) + pos_index = nEllipses - 1; + macroSizesEll_arr[pos_index] += m_size; + } + } + // calculates sum over all CPUs + ORBIT_MPI_Allreduce(macroSizesEll_arr, macroSizesEll_MPI_arr, nEllipses, MPI_DOUBLE, MPI_SUM, bunch->getMPI_Comm_Local()->comm); + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_arr[ie] = macroSizesEll_MPI_arr[ie]; + // std::cout<<"debug 0 ie ="<< ie <<" macrosize="<< macroSizesEll_MPI_arr[ie] << std::endl; + } + // calculate the relative volume density in each region. This density is a sum of all elipsoids + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_MPI_arr[ie] /= ((ie + 2) * (ie + 2) * (ie + 2) - (ie + 1) * (ie + 1) * (ie + 1)); + // std::cout<<"debug 1 ie ="<< ie <<" macrosize="<< macroSizesEll_MPI_arr[ie] << std::endl; + } + // calculate the density for each elipsoid + double rho_sum = 0.; + for (int ie = (nEllipses - 1); ie >= 0; ie--) { + macroSizesEll_MPI_arr[ie] -= rho_sum; + rho_sum += macroSizesEll_MPI_arr[ie]; + // std::cout<<"debug 2 ie ="<< ie <<" macrosize="<< macroSizesEll_MPI_arr[ie] << " rho_sum="<< rho_sum < 0) { + ParticleMacroSize *macroSizeAttr = (ParticleMacroSize *)bunch->getParticleAttributes("macrosize"); + for (int ip = 0, n = bunch->getSize(); ip < n; ip++) { + double m_size = macroSizeAttr->macrosize(ip); + coordArr = partArr[ip]; + coord_avg[0] += m_size * coordArr[0]; + coord_avg[1] += m_size * coordArr[2]; + coord_avg[2] += m_size * coordArr[4]; + coord_avg[3] += m_size * coordArr[0] * coordArr[0]; + coord_avg[4] += m_size * coordArr[2] * coordArr[2]; + coord_avg[5] += m_size * coordArr[4] * coordArr[4]; + coord_avg[6] += m_size * coordArr[0] * coordArr[2]; + coord_avg[7] += m_size; + } + } else { + double m_size = bunch->getMacroSize(); + int nParts = bunch->getSize(); + coord_avg[7] = m_size * nParts; + for (int ip = 0; ip < nParts; ip++) { + coordArr = partArr[ip]; + coord_avg[0] += coordArr[0]; + coord_avg[1] += coordArr[2]; + coord_avg[2] += coordArr[4]; + coord_avg[3] += coordArr[0] * coordArr[0]; + coord_avg[4] += coordArr[2] * coordArr[2]; + coord_avg[5] += coordArr[4] * coordArr[4]; + coord_avg[6] += coordArr[0] * coordArr[2]; + } + for (int i = 0; i < 7; i++) { + coord_avg[i] *= m_size; + } + } + + ORBIT_MPI_Allreduce(coord_avg, coord_avg_out, 8, MPI_DOUBLE, MPI_SUM, bunch->getMPI_Comm_Local()->comm); + + total_macrosize = coord_avg_out[7]; + if (total_macrosize == 0.) { + BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index0); + BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index1); + return; + } + + x_center = coord_avg_out[0] / total_macrosize; + y_center = coord_avg_out[1] / total_macrosize; + z_center = coord_avg_out[2] / total_macrosize; + double cov_xx = coord_avg_out[3] / total_macrosize - x_center * x_center; + double cov_yy = coord_avg_out[4] / total_macrosize - y_center * y_center; + double cov_xy = coord_avg_out[6] / total_macrosize - x_center * y_center; + z2_avg = fabs(coord_avg_out[5] / total_macrosize - z_center * z_center); + + double phi = -0.5 * atan2(2. * cov_xy, cov_xx - cov_yy); + cos_phi = cos(phi); + sin_phi = sin(phi); + double sin_cos_phi = sin_phi * cos_phi; + x2_avg = fabs(cov_xx * cos_phi * cos_phi + cov_yy * sin_phi * sin_phi - 2. * cov_xy * sin_cos_phi); + y2_avg = fabs(cov_xx * sin_phi * sin_phi + cov_yy * cos_phi * cos_phi + 2. * cov_xy * sin_cos_phi); + a2_ellips = 4. * x2_avg; + b2_ellips = 4. * y2_avg; + a_ellips = sqrt(a2_ellips); + b_ellips = sqrt(b2_ellips); + bunch_length = sqrt(12. * z2_avg); + + BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index0); + BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index1); + + if (x2_avg == 0. || y2_avg == 0. || z2_avg == 0.) { + int rank = 0; + ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) { + std::cerr << "SpaceChargeCalcUnifEllipse2D::bunchAnalysis(bunch,...)" << std::endl + << "The bunch rms sizes are wrong! Cannot calculate space charge!" << std::endl + << " x2_rms=" << x2_avg << std::endl + << " y2_rms=" << y2_avg << std::endl + << " z2_rms=" << z2_avg << std::endl + << "Stop." << std::endl; + } + ORBIT_MPI_Finalize(); + } + + if (nEllipses == 1) { + macroSizesEll_arr[0] = total_macrosize; + ellipsoidCalc_arr[0]->setEllipse(a_ellips, b_ellips); + ellipsoidCalc_arr[0]->setQ(total_macrosize / bunch_length); + return; + } + + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_arr[ie] = 0.; + } + + for (int ip = 0, n = bunch->getSize(); ip < n; ip++) { + coordArr = partArr[ip]; + double x = coordArr[0] - x_center; + double y = coordArr[2] - y_center; + double x_rot = cos_phi * x - sin_phi * y; + double y_rot = sin_phi * x + cos_phi * y; + double pos = sqrt(x_rot * x_rot / a2_ellips + y_rot * y_rot / b2_ellips); + int pos_index = int(pos * nEllipses) - 1; + if (pos_index < 0) + pos_index = 0; + if (pos_index >= nEllipses) + pos_index = nEllipses - 1; + double m_size = bunch->getMacroSize(); + if (has_msize > 0) { + ParticleMacroSize *macroSizeAttr = (ParticleMacroSize *)bunch->getParticleAttributes("macrosize"); + m_size = macroSizeAttr->macrosize(ip); + } + macroSizesEll_arr[pos_index] += m_size; + } + + ORBIT_MPI_Allreduce(macroSizesEll_arr, macroSizesEll_MPI_arr, nEllipses, MPI_DOUBLE, MPI_SUM, bunch->getMPI_Comm_Local()->comm); + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_arr[ie] = macroSizesEll_MPI_arr[ie]; + macroSizesEll_MPI_arr[ie] /= ((ie + 2) * (ie + 2) - (ie + 1) * (ie + 1)); + } + + double rho_sum = 0.; + for (int ie = nEllipses - 1; ie >= 0; ie--) { + macroSizesEll_MPI_arr[ie] -= rho_sum; + rho_sum += macroSizesEll_MPI_arr[ie]; + } + + double q_sum = 0.; + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_MPI_arr[ie] *= (ie + 1) * (ie + 1); + q_sum += macroSizesEll_MPI_arr[ie]; + } + double q_coeff = total_macrosize / q_sum; + for (int ie = 0; ie < nEllipses; ie++) { + macroSizesEll_arr[ie] = macroSizesEll_MPI_arr[ie] * q_coeff; + double coeff = (ie + 2.) / nEllipses; + ellipsoidCalc_arr[ie]->setEllipse(a_ellips * coeff, b_ellips * coeff); + ellipsoidCalc_arr[ie]->setQ(macroSizesEll_arr[ie] / bunch_length); + } +} + +/** Calculates the electric field in the bunch-centered system. */ +void SpaceChargeCalcUnifEllipse2D::calculateField(double x, double y, double z, double &ex, double &ey, double &ez) { + double x_rot = cos_phi * x - sin_phi * y; + double y_rot = sin_phi * x + cos_phi * y; + double x2 = x_rot * x_rot; + double y2 = y_rot * y_rot; + double ex_rot = 0.; + double ey_rot = 0.; + for (int ie = 0; ie < nEllipses; ie++) { + double ex_l, ey_l; + ellipsoidCalc_arr[ie]->calcField2D(x_rot, y_rot, x2, y2, ex_l, ey_l); + ex_rot += ex_l; + ey_rot += ey_l; + } + ex = cos_phi * ex_rot + sin_phi * ey_rot; + ey = -sin_phi * ex_rot + cos_phi * ey_rot; + ez = 0.; } diff --git a/src/spacecharge/SpaceChargeCalcUnifEllipse.hh b/src/spacecharge/SpaceChargeCalcUnifEllipse.hh index d88f9ba3..8de78eac 100644 --- a/src/spacecharge/SpaceChargeCalcUnifEllipse.hh +++ b/src/spacecharge/SpaceChargeCalcUnifEllipse.hh @@ -4,78 +4,106 @@ The space charge kick is transformed later into the lab system. */ -#ifndef SC_SPACECHARGE_CALC_UNIFORM_ELLIPSE_HH -#define SC_SPACECHARGE_CALC_UNIFORM_ELLIPSE_HH +#ifndef SPACE_CHARGE_CALC_UNIFORM_ELLIPSE_HH +#define SPACE_CHARGE_CALC_UNIFORM_ELLIPSE_HH -//MPI Function Wrappers +// MPI Function Wrappers #include "orbit_mpi.hh" #include "wrap_mpi_comm.hh" -#include #include +#include -//ORBIT bunch +// ORBIT bunch #include "Bunch.hh" -//pyORBIT utils +// pyORBIT utils #include "CppPyWrapper.hh" #include "UniformEllipsoidFieldCalculator.hh" using namespace std; -class SpaceChargeCalcUnifEllipse: public OrbitUtils::CppPyWrapper -{ -public: - - /** Constructor with the "x to y ratio" parameter. */ - SpaceChargeCalcUnifEllipse(int nEllipses_in); - - /** Destructor */ - virtual ~SpaceChargeCalcUnifEllipse(); - - /** Calculates space charge and applies the transverse and - longitudinal SC kicks to the macro-particles in the bunch. */ - void trackBunch(Bunch* bunch, double length); - - /** Analyses the bunch and sets up the ellipsoid filed sources */ - void bunchAnalysis(Bunch* bunch); - - /** Calculates the electric filed in the center of the bunch sytem. */ - void calculateField(double x, double y, double z, double& ex, double& ey, double& ez) ; - - /** Returns the UniformEllipsoidFieldCalculator class instance with a particular index */ - UniformEllipsoidFieldCalculator* getEllipsFieldCalculator(int ellipse_index); - - /** Returns the number of UniformEllipsoidFieldCalculator class instances */ - int getNEllipses(); - -private: - -protected: - - //number of uniform ellipses - int nEllipses; - - //total macrosize - double total_macrosize; +class SpaceChargeCalcUnifEllipse : public OrbitUtils::CppPyWrapper { + public: + /** Constructor with the "x to y ratio" parameter. */ + SpaceChargeCalcUnifEllipse(int nEllipses_in); + + /** Destructor. */ + virtual ~SpaceChargeCalcUnifEllipse(); + + /** Calculates space charge and applies the transverse and longitudinal SC kicks to the macro-particles in the bunch. */ + virtual void trackBunch(Bunch *bunch, double length); + + /** Analyzes the bunch and sets up the ellipsoid filed sources. */ + virtual void bunchAnalysis(Bunch *bunch); + + /** Calculates the electric filed in the center of the bunch system. */ + virtual void calculateField(double x, double y, double z, double &ex, double &ey, double &ez); + + /** Returns the UniformEllipsoidFieldCalculator class instance with a particular index. */ + UniformEllipsoidFieldCalculator *getEllipsFieldCalculator(int ellipse_index); + + /** Returns the number of UniformEllipsoidFieldCalculator class instances. */ + int getNEllipses(); + + private: + protected: + // Number of ellipsoids + int nEllipses; + + // Total macrosize + double total_macrosize; + + // Distribution parameters + double x_center; + double y_center; + double z_center; + double x2_avg; + double y2_avg; + double z2_avg; + double xMin; + double xMax; + double yMin; + double yMax; + double zMin; + double zMax; + + // Sizes of the biggest ellipsoid + double a_ellips; + double b_ellips; + double c_ellips; + double a2_ellips; + double b2_ellips; + double c2_ellips; + + // Field calculators + UniformEllipsoidFieldCalculator **ellipsoidCalc_arr; + + // Total macrosize in each ellipsoid + double *macroSizesEll_arr; + double *macroSizesEll_MPI_arr; +}; - //parameters of the distribution - double x_center, y_center, z_center; - double x2_avg, y2_avg, z2_avg; - double xMin, xMax, yMin, yMax, zMin, zMax; +class SpaceChargeCalcUnifEllipse2D : public SpaceChargeCalcUnifEllipse { + public: + /** Constructor. */ + SpaceChargeCalcUnifEllipse2D(int nEllipses_in); - //sizes of the biggest ellipsoid - double a_ellips, b_ellips, c_ellips; - double a2_ellips, b2_ellips, c2_ellips; + /** Calculates and applies transverse space-charge kicks. */ + void trackBunch(Bunch *bunch, double length) override; - //ellipse calculators - UniformEllipsoidFieldCalculator** ellipsoidCalc_arr; + /** Analyses the bunch and sets up the ellipse field sources. */ + void bunchAnalysis(Bunch *bunch) override; - //total macrosize in each ellipsoid - double* macroSizesEll_arr; - double* macroSizesEll_MPI_arr; + /** Calculates the electric field in the bunch-centered system. */ + void calculateField(double x, double y, double z, double &ex, double &ey, double &ez) override; + protected: + double cos_phi; + double sin_phi; + double bunch_length; }; -//end of SC_SPACECHARGE_CALC_UNIFORM_ELLIPSE_HH +// end of SPACE_CHARGE_CALC_UNIFORM_ELLIPSE_HH + #endif diff --git a/src/spacecharge/UniformEllipsoidFieldCalculator.cc b/src/spacecharge/UniformEllipsoidFieldCalculator.cc index 811eb52d..3cc41914 100644 --- a/src/spacecharge/UniformEllipsoidFieldCalculator.cc +++ b/src/spacecharge/UniformEllipsoidFieldCalculator.cc @@ -1,9 +1,12 @@ /** - This class calculates the field of uniformly charged ellipsoid by using - the symmetric elliptic integral and Carlson formulas for these integrals. - */ + Calculates the electric field generated by a uniformly charged ellipsoid in free space. -//MPI Function Wrappers + - 3D fields: calculated from symmetric elliptic integrals, evaluated using Carlson's formulas. + - 2D transverse fields (long-bunch approximation): calculated from analytic formulas. + - 1D longitudinal fields (long-bunch approximation): to do. + **/ + +// MPI Function Wrappers #include "orbit_mpi.hh" #include "wrap_mpi_comm.hh" @@ -14,268 +17,299 @@ using namespace OrbitUtils; -//macros for max and min +// macros for max and min #if !defined(DOXYGEN_SHOULD_SKIP_THIS) #ifndef max - #define max( a, b ) ( ((a) > (b)) ? (a) : (b) ) +#define max(a, b) (((a) > (b)) ? (a) : (b)) #endif #ifndef min - #define min( a, b ) ( ((a) < (b)) ? (a) : (b) ) +#define min(a, b) (((a) < (b)) ? (a) : (b)) #endif #endif -/** Constructor. There is no parameters */ -UniformEllipsoidFieldCalculator::UniformEllipsoidFieldCalculator(): CppPyWrapper(NULL) -{ - intFuncX0 = new Function(); - intFuncY0 = new Function(); - intFuncZ0 = new Function(); - intFuncX1 = new Function(); - intFuncY1 = new Function(); - intFuncZ1 = new Function(); - intFuncX2 = new Function(); - intFuncY2 = new Function(); - intFuncZ2 = new Function(); - //the number of points - lambda_function_points0 = 200; - - //the number of points - lambda_function_points1 = 50; +/** Constructor. */ +UniformEllipsoidFieldCalculator::UniformEllipsoidFieldCalculator() : CppPyWrapper(NULL) { + intFuncX0 = new Function(); + intFuncY0 = new Function(); + intFuncZ0 = new Function(); + intFuncX1 = new Function(); + intFuncY1 = new Function(); + intFuncZ1 = new Function(); + intFuncX2 = new Function(); + intFuncY2 = new Function(); + intFuncZ2 = new Function(); - //the number of points - lambda_function_points2 = 50; + // Number of points for lambda function + lambda_function_points0 = 200; + lambda_function_points1 = 50; + lambda_function_points2 = 50; - //Q_total is 1 by dfeault - Q_total = 1.0; + // Total charge + Q_total = 1.0; - //the parameter of ellipses - a = 1.; b = 1.; c = 1.; - double r_max = 10.; - setEllipsoid(a,b,c,r_max); + // Ellipse parameters (x/a)^2 + (y/b)^2 + (z/c)^2 = 1 + a = 1.0; + b = 1.0; + c = 1.0; + double r_max = 10.0; + setEllipsoid(a, b, c, r_max); } /** Destructor */ -UniformEllipsoidFieldCalculator::~UniformEllipsoidFieldCalculator() -{ - delete intFuncX0; - delete intFuncY0; - delete intFuncZ0; - delete intFuncX1; - delete intFuncY1; - delete intFuncZ1; - delete intFuncX2; - delete intFuncY2; - delete intFuncZ2; +UniformEllipsoidFieldCalculator::~UniformEllipsoidFieldCalculator() { + delete intFuncX0; + delete intFuncY0; + delete intFuncZ0; + delete intFuncX1; + delete intFuncY1; + delete intFuncZ1; + delete intFuncX2; + delete intFuncY2; + delete intFuncZ2; } +/** Sets ellipsoid semi-axes (a, b, c). */ +void UniformEllipsoidFieldCalculator::setEllipsoid(double a_in, double b_in, double c_in, double r_max) { + a = a_in; + b = b_in; + c = c_in; + a2 = a * a; + b2 = b * b; + c2 = c * c; -/** Sets the half-axis of the ellipsoid */ -void UniformEllipsoidFieldCalculator::setEllipsoid(double a_in, double b_in, double c_in, double r_max) -{ - a = a_in; b = b_in; c = c_in; - a2 = a*a; b2 = b*b; c2 = c*c; + lambda_max2 = pow(r_max, 2); + lambda_max0 = pow(2.0 * max(max(a, b), c), 2); + lambda_max1 = pow(5. * max(max(a, b), c), 2); + if (lambda_max2 < lambda_max0) { + lambda_max0 = lambda_max2; + lambda_max1 = lambda_max2; + intFuncX1->clean(); + intFuncY1->clean(); + intFuncZ1->clean(); + intFuncX2->clean(); + intFuncY2->clean(); + intFuncZ2->clean(); + } + if (lambda_max2 < lambda_max1) { + lambda_max1 = lambda_max2; + intFuncX2->clean(); + intFuncY2->clean(); + intFuncZ2->clean(); + } + // Interval from 0 to lambda_max0. + lambda_eps = 0.01 * lambda_max0 / (lambda_function_points0 - 1); + + // Integrate the part from lambda to lambda_max for different lambda and put values into functions. + double lambda_step = lambda_max0 / (lambda_function_points0 - 1); + intFuncX0->clean(); + intFuncY0->clean(); + intFuncZ0->clean(); + for (int iL = 0, nL = lambda_function_points0; iL < nL; iL++) { + double lambda = lambda_step * iL; + intFuncX0->add(lambda, this->integralPhi(a2, b2, c2, lambda)); + intFuncY0->add(lambda, this->integralPhi(b2, a2, c2, lambda)); + intFuncZ0->add(lambda, this->integralPhi(c2, b2, a2, lambda)); + } + intFuncX0->setConstStep(1); + intFuncY0->setConstStep(1); + intFuncZ0->setConstStep(1); + if (intFuncX0->isStepConst() != 1 || intFuncY0->isStepConst() != 1 || intFuncZ0->isStepConst() != 1) { + int rank = 0; + ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) { + std::cerr << "UniformEllipsoidFieldCalculator::setEllipsoid(...)" << std::endl + << "The Functions for lambda are not equidistant!!! " << std::endl + << "Stop." << std::endl; + } + ORBIT_MPI_Finalize(); + } + if (lambda_max1 > lambda_max0) { + lambda_step = (lambda_max1 - lambda_max0) / (lambda_function_points1 - 1); + intFuncX1->clean(); + intFuncY1->clean(); + intFuncZ1->clean(); + for (int iL = 0, nL = lambda_function_points1; iL < nL; iL++) { + double lambda = lambda_max0 + lambda_step * iL; + intFuncX1->add(lambda, this->integralPhi(a2, b2, c2, lambda) * lambda); + intFuncY1->add(lambda, this->integralPhi(b2, a2, c2, lambda) * lambda); + intFuncZ1->add(lambda, this->integralPhi(c2, b2, a2, lambda) * lambda); + } + intFuncX1->setConstStep(1); + intFuncY1->setConstStep(1); + intFuncZ1->setConstStep(1); + if (intFuncX1->isStepConst() != 1 || intFuncY1->isStepConst() != 1 || intFuncZ1->isStepConst() != 1) { + int rank = 0; + ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) { + std::cerr << "UniformEllipsoidFieldCalculator::setEllipsoid(...)" << std::endl + << "The Functions for lambda are not equidistant!!! " << std::endl + << "Stop." << std::endl; + } + ORBIT_MPI_Finalize(); + } + } + if (lambda_max2 > lambda_max1) { + lambda_step = (lambda_max2 - lambda_max1) / (lambda_function_points2 - 1); + intFuncX2->clean(); + intFuncY2->clean(); + intFuncZ2->clean(); + for (int iL = 0, nL = lambda_function_points2; iL < nL; iL++) { + double lambda = lambda_max1 + lambda_step * iL; + intFuncX2->add(lambda, this->integralPhi(a2, b2, c2, lambda) * lambda); + intFuncY2->add(lambda, this->integralPhi(b2, a2, c2, lambda) * lambda); + intFuncZ2->add(lambda, this->integralPhi(c2, b2, a2, lambda) * lambda); + } + intFuncX2->setConstStep(1); + intFuncY2->setConstStep(1); + intFuncZ2->setConstStep(1); + if (intFuncX2->isStepConst() != 1 || intFuncY2->isStepConst() != 1 || intFuncZ2->isStepConst() != 1) { + int rank = 0; + ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) { + std::cerr << "UniformEllipsoidFieldCalculator::setEllipsoid(...)" << std::endl + << "The Functions for lambda are not equidistant!!! " << std::endl + << "Stop." << std::endl; + } + ORBIT_MPI_Finalize(); + } + } +} - lambda_max2 = pow(r_max,2); - lambda_max0 = pow(2.0*max(max(a,b),c),2); - lambda_max1 = pow(5.*max(max(a,b),c),2); - if(lambda_max2 < lambda_max0){ - lambda_max0 = lambda_max2; - lambda_max1 = lambda_max2; - intFuncX1->clean(); - intFuncY1->clean(); - intFuncZ1->clean(); - intFuncX2->clean(); - intFuncY2->clean(); - intFuncZ2->clean(); - } - if(lambda_max2 < lambda_max1){ - lambda_max1 = lambda_max2; - intFuncX2->clean(); - intFuncY2->clean(); - intFuncZ2->clean(); - } - //interval from 0 to lambda_max0 - lambda_eps = 0.01*lambda_max0/(lambda_function_points0 - 1); - //now integrate the part from lambda to lambda_max for different lambda and put values into functions - double lambda_step = lambda_max0/(lambda_function_points0 - 1); - intFuncX0->clean(); - intFuncY0->clean(); - intFuncZ0->clean(); - for(int iL = 0, nL = lambda_function_points0; iL < nL; iL++){ - double lambda = lambda_step*iL; - intFuncX0->add(lambda, this->integralPhi(a2,b2,c2,lambda)); - intFuncY0->add(lambda, this->integralPhi(b2,a2,c2,lambda)); - intFuncZ0->add(lambda, this->integralPhi(c2,b2,a2,lambda)); - } - intFuncX0->setConstStep(1); - intFuncY0->setConstStep(1); - intFuncZ0->setConstStep(1); - if(intFuncX0->isStepConst() != 1 || intFuncY0->isStepConst() != 1 || intFuncZ0->isStepConst() != 1){ - int rank = 0; - ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); - if(rank == 0){ - std::cerr << "UniformEllipsoidFieldCalculator::setEllipsoid(...)" << std::endl - << "The Functions for lambda are not equidistant!!! "<< std::endl - << "Stop."<< std::endl; - } - ORBIT_MPI_Finalize(); - } - if(lambda_max1 > lambda_max0){ - lambda_step = (lambda_max1-lambda_max0)/(lambda_function_points1 - 1); - intFuncX1->clean(); - intFuncY1->clean(); - intFuncZ1->clean(); - for(int iL = 0, nL = lambda_function_points1; iL < nL; iL++){ - double lambda = lambda_max0 + lambda_step*iL; - intFuncX1->add(lambda, this->integralPhi(a2,b2,c2,lambda)*lambda); - intFuncY1->add(lambda, this->integralPhi(b2,a2,c2,lambda)*lambda); - intFuncZ1->add(lambda, this->integralPhi(c2,b2,a2,lambda)*lambda); - } - intFuncX1->setConstStep(1); - intFuncY1->setConstStep(1); - intFuncZ1->setConstStep(1); - if(intFuncX1->isStepConst() != 1 || intFuncY1->isStepConst() != 1 || intFuncZ1->isStepConst() != 1){ - int rank = 0; - ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); - if(rank == 0){ - std::cerr << "UniformEllipsoidFieldCalculator::setEllipsoid(...)" << std::endl - << "The Functions for lambda are not equidistant!!! "<< std::endl - << "Stop."<< std::endl; - } - ORBIT_MPI_Finalize(); - } - } - if(lambda_max2 > lambda_max1){ - lambda_step = (lambda_max2 - lambda_max1)/(lambda_function_points2 - 1); - intFuncX2->clean(); - intFuncY2->clean(); - intFuncZ2->clean(); - for(int iL = 0, nL = lambda_function_points2; iL < nL; iL++){ - double lambda = lambda_max1 + lambda_step*iL; - intFuncX2->add(lambda, this->integralPhi(a2,b2,c2,lambda)*lambda); - intFuncY2->add(lambda, this->integralPhi(b2,a2,c2,lambda)*lambda); - intFuncZ2->add(lambda, this->integralPhi(c2,b2,a2,lambda)*lambda); - } - intFuncX2->setConstStep(1); - intFuncY2->setConstStep(1); - intFuncZ2->setConstStep(1); - if(intFuncX2->isStepConst() != 1 || intFuncY2->isStepConst() != 1 || intFuncZ2->isStepConst() != 1){ - int rank = 0; - ORBIT_MPI_Comm_rank(MPI_COMM_WORLD, &rank); - if(rank == 0){ - std::cerr << "UniformEllipsoidFieldCalculator::setEllipsoid(...)" << std::endl - << "The Functions for lambda are not equidistant!!! "<< std::endl - << "Stop."<< std::endl; - } - ORBIT_MPI_Finalize(); - } - } +/** Sets ellipsoid semi-axes in transverse plane (a, b). */ +void UniformEllipsoidFieldCalculator::setEllipse(double a_in, double b_in) { + a = a_in; + b = b_in; + a2 = a * a; + b2 = b * b; } /** Calculates the field components */ -void UniformEllipsoidFieldCalculator::calcField(double x, double y, double z, - double x2, double y2, double z2, - double& ex, double& ey, double& ez) -{ - if((x2/a2+y2/b2+z2/c2) <= 1.){ - ex = Q_total*x*intFuncX0->y(0); - ey = Q_total*y*intFuncY0->y(0); - ez = Q_total*z*intFuncZ0->y(0); - return; - } +void UniformEllipsoidFieldCalculator::calcField(double x, double y, double z, double x2, double y2, double z2, double &ex, double &ey, double &ez) { + if ((x2 / a2 + y2 / b2 + z2 / c2) <= 1.) { + ex = Q_total * x * intFuncX0->y(0); + ey = Q_total * y * intFuncY0->y(0); + ez = Q_total * z * intFuncZ0->y(0); + return; + } + + double lambda = this->calcLambda(x, y, z, x2, y2, z2); + if (lambda < lambda_max0) { + ex = Q_total * x * intFuncX0->getY(lambda); + ey = Q_total * y * intFuncY0->getY(lambda); + ez = Q_total * z * intFuncZ0->getY(lambda); + return; + } else if (lambda < lambda_max1) { + ex = Q_total * x * intFuncX1->getY(lambda) / lambda; + ey = Q_total * y * intFuncY1->getY(lambda) / lambda; + ez = Q_total * z * intFuncZ1->getY(lambda) / lambda; + return; + } else if (lambda < lambda_max2) { + ex = Q_total * x * intFuncX2->getY(lambda) / lambda; + ey = Q_total * y * intFuncY2->getY(lambda) / lambda; + ez = Q_total * z * intFuncZ2->getY(lambda) / lambda; + return; + } + double r2 = x2 + y2 + z2; + double r3 = sqrt(r2) * r2; + ex = Q_total * x / r3; + ey = Q_total * y / r3; + ez = Q_total * z / r3; +} + +/** Calculates the field components of a uniformly charged ellipsoid. */ +void UniformEllipsoidFieldCalculator::calcField2D(double x, double y, double x2, double y2, double &ex, double &ey) { + double lambda = 0.; + if ((x2 / a2 + y2 / b2) > 1.) { + double p = a2 + b2 - x2 - y2; + double q = a2 * b2 - x2 * b2 - y2 * a2; + double discriminant = p * p - 4. * q; + if (discriminant < 0.) + discriminant = 0.; + lambda = 0.5 * (-p + sqrt(discriminant)); + } - double lambda = this->calcLambda(x,y,z,x2,y2,z2); - if(lambda < lambda_max0){ - ex = Q_total*x*intFuncX0->getY(lambda); - ey = Q_total*y*intFuncY0->getY(lambda); - ez = Q_total*z*intFuncZ0->getY(lambda); - return; - } else if(lambda < lambda_max1) { - ex = Q_total*x*intFuncX1->getY(lambda)/lambda; - ey = Q_total*y*intFuncY1->getY(lambda)/lambda; - ez = Q_total*z*intFuncZ1->getY(lambda)/lambda; - return; - } else if(lambda < lambda_max2) { - ex = Q_total*x*intFuncX2->getY(lambda)/lambda; - ey = Q_total*y*intFuncY2->getY(lambda)/lambda; - ez = Q_total*z*intFuncZ2->getY(lambda)/lambda; - return; - } - double r2 = x2+y2+z2; - double r3 = sqrt(r2)*r2; - ex = Q_total*x/r3; ey = Q_total*y/r3; ez = Q_total*z/r3; + double a_lambda = sqrt(a2 + lambda); + double b_lambda = sqrt(b2 + lambda); + ex = 2. * Q_total * x / (a_lambda * (a_lambda + b_lambda)); + ey = 2. * Q_total * y / (b_lambda * (a_lambda + b_lambda)); } -/** Calculates lambda value as a root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 */ -double UniformEllipsoidFieldCalculator::calcLambda(double x, double y, double z, - double x2, double y2, double z2) -{ - double lambda_start = 0.; - double a2_ls, b2_ls, c2_ls; - a2_ls = a2+lambda_start; b2_ls = b2+lambda_start; c2_ls = c2+lambda_start; - double v_start = x2/a2_ls + y2/b2_ls + z2/c2_ls - 1.0; - double lambda_stop = lambda_max2; - double v_stop = x2/(a2+lambda_stop) + y2/(b2+lambda_stop) + z2/(c2+lambda_stop) - 1.0; - double vp_start = x2/pow((a2+lambda_start),2) + y2/pow((b2+lambda_start),2) + z2/pow((c2+lambda_start),2); - lambda_stop = lambda_start - v_start*(lambda_start - lambda_stop)/(v_start - v_stop); - lambda_start = lambda_start + v_start/vp_start; - while(fabs(lambda_start - lambda_stop) > lambda_eps){ - //std::cout << "debug calcFieldlambda_start ="< #include +#include using namespace std; /** This class calculates the field of uniformly charged ellipsoid by using - the symmetric elliptic integral and Carlson formulas for these integrals. + the symmetric elliptic integral and Carlson formulas for these integrals. */ - -class UniformEllipsoidFieldCalculator: public OrbitUtils::CppPyWrapper -{ +class UniformEllipsoidFieldCalculator : public OrbitUtils::CppPyWrapper { public: + /** Constructor */ + UniformEllipsoidFieldCalculator(); - /** Constructor */ - UniformEllipsoidFieldCalculator(); + /** Destructor */ + virtual ~UniformEllipsoidFieldCalculator(); - /** Destructor */ - virtual ~UniformEllipsoidFieldCalculator(); + /** Sets the half-axis of the ellipsoid and maximal values of radius */ + void setEllipsoid(double a_in, double b_in, double c_in, double r_max); - /** Sets the half-axis of the ellipsoid and maximal values of radius */ - void setEllipsoid(double a_in, double b_in, double c_in, double r_max); + /** Sets the half-axes of the ellipse. */ + void setEllipse(double a_in, double b_in); - /** Calculates the field components */ - void calcField(double x, double y, double z, - double x2, double y2, double z2, - double& ex, double& ey, double& ez); + /** Calculates the field components */ + void calcField(double x, double y, double z, + double x2, double y2, double z2, + double &ex, double &ey, double &ez); - /** Calculates lambda value as a root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 */ - double calcLambda(double x, double y, double z, - double x2, double y2, double z2) ; + /** Calculates the field components of a uniformly charged ellipse. */ + void calcField2D(double x, double y, double x2, double y2, + double &ex, double &ey); - /** Returns the total space charge inside the ellipse. */ - double getQ(); + /** Calculates lambda value as a root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 */ + double calcLambda(double x, double y, double z, + double x2, double y2, double z2); - /** Sets the total space charge inside the ellipse. */ - void setQ(double Q_in); + /** Returns the total space charge inside the ellipse. */ + double getQ(); - private: + /** Sets the total space charge inside the ellipse. */ + void setQ(double Q_in); - /** Calculates integral for int(1.5*(1/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity */ - static double integralPhi(double a_2, double b_2, double c_2, double lambda); - - private: - - //total charge Q in the units of the electron cahrge - double Q_total; - - //the half-axis of the ellipsoid - double a,b,c; - double a2,b2,c2; - - //accuracy for the root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 - double lambda_eps; - - //lambda is a value for root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 - //number of points in the function of lambda for interval 0-2.5*max(a,b,c) - int lambda_function_points0; - double lambda_max0; - //number of points in the function of lambda for interval 2.5-5*max(a,b,c) - int lambda_function_points1; - double lambda_max1; - //number of points in the function of lambda for interval 5-lambda_max - int lambda_function_points2; - double lambda_max2; - - //values of integral for x,y,z axis lambda for interval 0 to 0-2.0*max(a,b,c) - // for x int((2/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - // for y int((2/(b^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - // for z int((2/(c^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - OrbitUtils::Function* intFuncX0; - OrbitUtils::Function* intFuncY0; - OrbitUtils::Function* intFuncZ0; - - //values of integral for x,y,z axis lambda for interval 2.5-5*max(a,b,c) - // for x int((2/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - // for y int((2/(b^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - // for z int((2/(c^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - OrbitUtils::Function* intFuncX1; - OrbitUtils::Function* intFuncY1; - OrbitUtils::Function* intFuncZ1; - - //values of integral for x,y,z axis lambda for interval 5-lambda_max - // for x int((2/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - // for y int((2/(b^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - // for z int((2/(c^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity - OrbitUtils::Function* intFuncX2; - OrbitUtils::Function* intFuncY2; - OrbitUtils::Function* intFuncZ2; + private: + /** Calculates integral for int(1.5*(1/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity */ + static double integralPhi(double a_2, double b_2, double c_2, double lambda); + private: + // total charge Q in the units of the electron cahrge + double Q_total; + + // the half-axis of the ellipsoid + double a, b, c; + double a2, b2, c2; + + // accuracy for the root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 + double lambda_eps; + + // lambda is a value for root of eq. x^2/(a^2+s) + y^2/(b^2+s) + z^2/(c^2+s) - 1 = 0 + // number of points in the function of lambda for interval 0-2.5*max(a,b,c) + int lambda_function_points0; + double lambda_max0; + // number of points in the function of lambda for interval 2.5-5*max(a,b,c) + int lambda_function_points1; + double lambda_max1; + // number of points in the function of lambda for interval 5-lambda_max + int lambda_function_points2; + double lambda_max2; + + // values of integral for x,y,z axis lambda for interval 0 to 0-2.0*max(a,b,c) + // for x int((2/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + // for y int((2/(b^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + // for z int((2/(c^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + OrbitUtils::Function *intFuncX0; + OrbitUtils::Function *intFuncY0; + OrbitUtils::Function *intFuncZ0; + + // values of integral for x,y,z axis lambda for interval 2.5-5*max(a,b,c) + // for x int((2/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + // for y int((2/(b^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + // for z int((2/(c^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + OrbitUtils::Function *intFuncX1; + OrbitUtils::Function *intFuncY1; + OrbitUtils::Function *intFuncZ1; + + // values of integral for x,y,z axis lambda for interval 5-lambda_max + // for x int((2/(a^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + // for y int((2/(b^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + // for z int((2/(c^2+s))*1/sqrt((a^2+s)*(b^2+s)*(c^2+s)), over s from lambda to infinity + OrbitUtils::Function *intFuncX2; + OrbitUtils::Function *intFuncY2; + OrbitUtils::Function *intFuncZ2; }; #endif diff --git a/src/spacecharge/wrap_spacecharge.cc b/src/spacecharge/wrap_spacecharge.cc index ab8d9ed5..8ec2e57e 100644 --- a/src/spacecharge/wrap_spacecharge.cc +++ b/src/spacecharge/wrap_spacecharge.cc @@ -15,7 +15,7 @@ #include "wrap_lspacechargecalc.hh" #include "wrap_spacechargecalc3d.hh" #include "wrap_uniform_ellipsoid_field_calculator.hh" -#include "wrap_spacechargecalc_uniform_ellipse.hh" +#include "wrap_spacechargecalc_unif_ellipse.hh" static PyMethodDef spacechargeMethods[] = { {NULL,NULL} }; @@ -39,7 +39,7 @@ extern "C" { wrap_spacecharge::initGrid2D(module); wrap_spacecharge::initGrid3D(module); wrap_spacecharge::initUniformEllipsoidFieldCalculator(module); - wrap_spacecharge::initSpaceChargeCalcUniformEllipse(module); + wrap_spacecharge::initSpaceChargeCalcUnifEllipse(module); wrap_spacecharge::initPoissonSolverFFT2D(module); wrap_spacecharge::initPoissonSolverFFT3D(module); wrap_spacecharge::initBoundary2D(module); diff --git a/src/spacecharge/wrap_spacechargecalc_unif_ellipse.cc b/src/spacecharge/wrap_spacechargecalc_unif_ellipse.cc new file mode 100644 index 00000000..b2e8e429 --- /dev/null +++ b/src/spacecharge/wrap_spacechargecalc_unif_ellipse.cc @@ -0,0 +1,247 @@ +#include "orbit_mpi.hh" +#include "pyORBIT_Object.hh" + +#include "wrap_bunch.hh" +#include "wrap_spacecharge.hh" +#include "wrap_spacechargecalc_unif_ellipse.hh" + +#include + +#include "SpaceChargeCalcUnifEllipse.hh" + +using namespace OrbitUtils; + +namespace wrap_spacecharge { + +#ifdef __cplusplus +extern "C" { +#endif + +//--------------------------------------------------------- +// Python SpaceChargeCalcUnifEllipse class definition +//--------------------------------------------------------- + +// constructor for python class wrapping SpaceChargeCalcUnifEllipse instance +// It never will be called directly + +static PyObject *SpaceChargeCalcUnifEllipse_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { + pyORBIT_Object *self; + self = (pyORBIT_Object *)type->tp_alloc(type, 0); + self->cpp_obj = NULL; + return (PyObject *)self; +} + +// initializator for python SpaceChargeCalcUnifEllipse class +// this is implementation of the __init__ method SpaceChargeCalcUnifEllipse(nEllipses = 1]) +static int SpaceChargeCalcUnifEllipse_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds) { + int nEllipses = 1; + if (!PyArg_ParseTuple(args, "|i:__init__", &nEllipses)) { + ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse - SpaceChargeCalcUnifEllipse([nEllipses = 1]) - constructor needs parameters."); + } + self->cpp_obj = new SpaceChargeCalcUnifEllipse(nEllipses); + return 0; +} + +// initializator for python SpaceChargeCalcUnifEllipse2D class +static int SpaceChargeCalcUnifEllipse2D_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds) { + int nEllipses = 1; + if (!PyArg_ParseTuple(args, "|i:__init__", &nEllipses)) { + ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse2D - SpaceChargeCalcUnifEllipse2D([nEllipses = 1]) - constructor needs parameters."); + } + self->cpp_obj = new SpaceChargeCalcUnifEllipse2D(nEllipses); + return 0; +} + +// trackBunch(Bunch* bunch, double length) +static PyObject *SpaceChargeCalcUnifEllipse_trackBunch(PyObject *self, PyObject *args) { + pyORBIT_Object *pySpaceChargeCalcUnifEllipse = (pyORBIT_Object *)self; + SpaceChargeCalcUnifEllipse *cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse *)pySpaceChargeCalcUnifEllipse->cpp_obj; + PyObject *pyBunch; + double length; + if (!PyArg_ParseTuple(args, "Od:trackBunch", &pyBunch, &length)) { + ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.trackBunch(pyBunch,length) - method needs parameters."); + } + PyObject *pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); + if (!PyObject_IsInstance(pyBunch, pyORBIT_Bunch_Type)) { + ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.trackBunch(pyBunch,length) - pyBunch is not Bunch."); + } + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)pyBunch)->cpp_obj; + cpp_SpaceChargeCalcUnifEllipse->trackBunch(cpp_bunch, length); + Py_INCREF(Py_None); + return Py_None; +} + +// UniformEllipsoidFieldCalculator* EllipsFieldCalculatorget(ellipse_index) returns the ellipse field object +static PyObject *SpaceChargeCalcUnifEllipse_getEllipsFieldCalculator(PyObject *self, PyObject *args) { + pyORBIT_Object *pySpaceChargeCalcUnifEllipse = (pyORBIT_Object *)self; + SpaceChargeCalcUnifEllipse *cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse *)pySpaceChargeCalcUnifEllipse->cpp_obj; + int index; + if (!PyArg_ParseTuple(args, "i:getEllipsFieldCalculator", &index)) { + ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.getEllipsFieldCalculator([index=0]) - method needs parameter."); + } + UniformEllipsoidFieldCalculator *cpp_ellipseFieldCalc = cpp_SpaceChargeCalcUnifEllipse->getEllipsFieldCalculator(index); + if (cpp_ellipseFieldCalc == NULL) { + Py_INCREF(Py_None); + return Py_None; + } + if (cpp_ellipseFieldCalc->getPyWrapper() != NULL) { + Py_INCREF(cpp_ellipseFieldCalc->getPyWrapper()); + return cpp_ellipseFieldCalc->getPyWrapper(); + } + // It will create a pyUniformEllipsoidFieldCalculator object + PyObject *mod = PyImport_ImportModule("orbit.core.spacecharge"); + PyObject *pyUniformEllipsoidFieldCalculator = PyObject_CallMethod(mod, const_cast("UniformEllipsoidFieldCalculator"), const_cast("")); + // delete the c++ reference to the internal UniformEllipsoidFieldCalculator inside pyUniformEllipsoidFieldCalculator and assign the new one + delete ((UniformEllipsoidFieldCalculator *)((pyORBIT_Object *)pyUniformEllipsoidFieldCalculator)->cpp_obj); + ((pyORBIT_Object *)pyUniformEllipsoidFieldCalculator)->cpp_obj = cpp_ellipseFieldCalc; + cpp_ellipseFieldCalc->setPyWrapper(pyUniformEllipsoidFieldCalculator); + Py_INCREF(cpp_ellipseFieldCalc->getPyWrapper()); + Py_DECREF(mod); + return pyUniformEllipsoidFieldCalculator; +} + +// getNEllipses() - returns the number of ellipses inside the Space Charge calculator +static PyObject *SpaceChargeCalcUnifEllipse_getNEllipses(PyObject *self, PyObject *args) { + pyORBIT_Object *pySpaceChargeCalcUnifEllipse = (pyORBIT_Object *)self; + SpaceChargeCalcUnifEllipse *cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse *)pySpaceChargeCalcUnifEllipse->cpp_obj; + return Py_BuildValue("i", cpp_SpaceChargeCalcUnifEllipse->getNEllipses()); +} + +// calculateField(x,y,z) - calculates the fileds from all ellipses +static PyObject *SpaceChargeCalcUnifEllipse_calculateField(PyObject *self, PyObject *args) { + pyORBIT_Object *pySpaceChargeCalcUnifEllipse = (pyORBIT_Object *)self; + SpaceChargeCalcUnifEllipse *cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse *)pySpaceChargeCalcUnifEllipse->cpp_obj; + double x, y, z, ex, ey, ez; + if (!PyArg_ParseTuple(args, "ddd:calculateField", &x, &y, &z)) { + ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.calculateField(x,y,z) - method needs parameters."); + } + cpp_SpaceChargeCalcUnifEllipse->calculateField(x, y, z, ex, ey, ez); + return Py_BuildValue("(ddd)", ex, ey, ez); + ; +} + +//----------------------------------------------------- +// destructor for python SpaceChargeCalcUnifEllipse class (__del__ method). +//----------------------------------------------------- +static void SpaceChargeCalcUnifEllipse_del(pyORBIT_Object *self) { + SpaceChargeCalcUnifEllipse *cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse *)self->cpp_obj; + if (cpp_SpaceChargeCalcUnifEllipse != NULL) { + delete cpp_SpaceChargeCalcUnifEllipse; + } + self->ob_base.ob_type->tp_free((PyObject *)self); +} + +// defenition of the methods of the python SpaceChargeCalcUnifEllipse wrapper class +// they will be vailable from python level +static PyMethodDef SpaceChargeCalcUnifEllipseClassMethods[] = { + {"trackBunch", SpaceChargeCalcUnifEllipse_trackBunch, METH_VARARGS, "track the bunch - trackBunch(pyBunch,length)"}, + {"getNEllipses", SpaceChargeCalcUnifEllipse_getNEllipses, METH_VARARGS, "returns the number of ellipses inside the Space Charge calculator"}, + {"getEllipsFieldCalculator", SpaceChargeCalcUnifEllipse_getEllipsFieldCalculator, METH_VARARGS, "returns the ellipse field object with particular index"}, + {"calculateField", SpaceChargeCalcUnifEllipse_calculateField, METH_VARARGS, "calculates the fileds ex,ey,ez from all ellipses"}, + {NULL}}; + +// defenition of the memebers of the python SpaceChargeCalcUnifEllipse wrapper class +// they will be vailable from python level +static PyMemberDef SpaceChargeCalcUnifEllipseClassMembers[] = { + {NULL}}; + +// new python SpaceChargeCalcUnifEllipse wrapper type definition +static PyTypeObject pyORBIT_SpaceChargeCalcUnifEllipse_Type = { + PyVarObject_HEAD_INIT(NULL, 0) "SpaceChargeCalcUnifEllipse", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor)SpaceChargeCalcUnifEllipse_del, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ + "The SpaceChargeCalcUnifEllipse python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + SpaceChargeCalcUnifEllipseClassMethods, /* tp_methods */ + SpaceChargeCalcUnifEllipseClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)SpaceChargeCalcUnifEllipse_init, /* tp_init */ + 0, /* tp_alloc */ + SpaceChargeCalcUnifEllipse_new, /* tp_new */ +}; + +static PyTypeObject pyORBIT_SpaceChargeCalcUnifEllipse2D_Type = { + PyVarObject_HEAD_INIT(NULL, 0) "SpaceChargeCalcUnifEllipse2D", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor)SpaceChargeCalcUnifEllipse_del, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ + "The SpaceChargeCalcUnifEllipse2D python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + SpaceChargeCalcUnifEllipseClassMethods, /* tp_methods */ + SpaceChargeCalcUnifEllipseClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)SpaceChargeCalcUnifEllipse2D_init, /* tp_init */ + 0, /* tp_alloc */ + SpaceChargeCalcUnifEllipse_new, /* tp_new */ +}; + +// Initialization function of the pySpaceChargeCalcUnifEllipse class. +// It will be called from SpaceCharge wrapper initialization. +void initSpaceChargeCalcUnifEllipse(PyObject *module) { + if (PyType_Ready(&pyORBIT_SpaceChargeCalcUnifEllipse_Type) < 0) + return; + Py_INCREF(&pyORBIT_SpaceChargeCalcUnifEllipse_Type); + PyModule_AddObject(module, "SpaceChargeCalcUnifEllipse", (PyObject *)&pyORBIT_SpaceChargeCalcUnifEllipse_Type); + if (PyType_Ready(&pyORBIT_SpaceChargeCalcUnifEllipse2D_Type) < 0) + return; + Py_INCREF(&pyORBIT_SpaceChargeCalcUnifEllipse2D_Type); + PyModule_AddObject(module, "SpaceChargeCalcUnifEllipse2D", (PyObject *)&pyORBIT_SpaceChargeCalcUnifEllipse2D_Type); +} + +#ifdef __cplusplus +} +#endif + +// end of namespace wrap_spacecharge +} // namespace wrap_spacecharge diff --git a/src/spacecharge/wrap_spacechargecalc_uniform_ellipse.hh b/src/spacecharge/wrap_spacechargecalc_unif_ellipse.hh similarity index 72% rename from src/spacecharge/wrap_spacechargecalc_uniform_ellipse.hh rename to src/spacecharge/wrap_spacechargecalc_unif_ellipse.hh index 85cc0fd9..dd627973 100644 --- a/src/spacecharge/wrap_spacechargecalc_uniform_ellipse.hh +++ b/src/spacecharge/wrap_spacechargecalc_unif_ellipse.hh @@ -7,9 +7,9 @@ extern "C" { #endif - namespace wrap_spacecharge{ - void initSpaceChargeCalcUniformEllipse(PyObject* module); - } +namespace wrap_spacecharge { +void initSpaceChargeCalcUnifEllipse(PyObject *module); +} #ifdef __cplusplus } diff --git a/src/spacecharge/wrap_spacechargecalc_uniform_ellipse.cc b/src/spacecharge/wrap_spacechargecalc_uniform_ellipse.cc deleted file mode 100644 index 657c610f..00000000 --- a/src/spacecharge/wrap_spacechargecalc_uniform_ellipse.cc +++ /dev/null @@ -1,197 +0,0 @@ -#include "orbit_mpi.hh" -#include "pyORBIT_Object.hh" - -#include "wrap_spacechargecalc_uniform_ellipse.hh" -#include "wrap_spacecharge.hh" -#include "wrap_bunch.hh" - -#include - -#include "SpaceChargeCalcUnifEllipse.hh" - -using namespace OrbitUtils; - -namespace wrap_spacecharge{ - -#ifdef __cplusplus -extern "C" { -#endif - - //--------------------------------------------------------- - //Python SpaceChargeCalcUnifEllipse class definition - //--------------------------------------------------------- - - //constructor for python class wrapping SpaceChargeCalcUnifEllipse instance - //It never will be called directly - - static PyObject* SpaceChargeCalcUnifEllipse_new(PyTypeObject *type, PyObject *args, PyObject *kwds) - { - pyORBIT_Object* self; - self = (pyORBIT_Object *) type->tp_alloc(type, 0); - self->cpp_obj = NULL; - return (PyObject *) self; - } - - //initializator for python SpaceChargeCalcUnifEllipse class - //this is implementation of the __init__ method SpaceChargeCalcUnifEllipse(nEllipses = 1]) - static int SpaceChargeCalcUnifEllipse_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ - int nEllipses = 1; - if(!PyArg_ParseTuple(args,"|i:__init__",&nEllipses)){ - ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse - SpaceChargeCalcUnifEllipse([nEllipses = 1]) - constructor needs parameters."); - } - self->cpp_obj = new SpaceChargeCalcUnifEllipse(nEllipses); - return 0; - } - - //trackBunch(Bunch* bunch, double length) - static PyObject* SpaceChargeCalcUnifEllipse_trackBunch(PyObject *self, PyObject *args){ - pyORBIT_Object* pySpaceChargeCalcUnifEllipse = (pyORBIT_Object*) self; - SpaceChargeCalcUnifEllipse* cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse*) pySpaceChargeCalcUnifEllipse->cpp_obj; - PyObject* pyBunch; - double length; - if(!PyArg_ParseTuple(args,"Od:trackBunch",&pyBunch,&length)){ - ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.trackBunch(pyBunch,length) - method needs parameters."); - } - PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); - if(!PyObject_IsInstance(pyBunch,pyORBIT_Bunch_Type)){ - ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.trackBunch(pyBunch,length) - pyBunch is not Bunch."); - } - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object*)pyBunch)->cpp_obj; - cpp_SpaceChargeCalcUnifEllipse->trackBunch(cpp_bunch,length); - Py_INCREF(Py_None); - return Py_None; - } - - //UniformEllipsoidFieldCalculator* EllipsFieldCalculatorget(ellipse_index) returns the ellipse field object - static PyObject* SpaceChargeCalcUnifEllipse_getEllipsFieldCalculator(PyObject *self, PyObject *args){ - pyORBIT_Object* pySpaceChargeCalcUnifEllipse = (pyORBIT_Object*) self; - SpaceChargeCalcUnifEllipse* cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse*) pySpaceChargeCalcUnifEllipse->cpp_obj; - int index; - if(!PyArg_ParseTuple(args,"i:getEllipsFieldCalculator",&index)){ - ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.getEllipsFieldCalculator([index=0]) - method needs parameter."); - } - UniformEllipsoidFieldCalculator* cpp_ellipseFieldCalc = cpp_SpaceChargeCalcUnifEllipse->getEllipsFieldCalculator(index); - if(cpp_ellipseFieldCalc == NULL) { - Py_INCREF(Py_None); - return Py_None; - } - if(cpp_ellipseFieldCalc->getPyWrapper() != NULL){ - Py_INCREF(cpp_ellipseFieldCalc->getPyWrapper()); - return cpp_ellipseFieldCalc->getPyWrapper(); - } - //It will create a pyUniformEllipsoidFieldCalculator object - PyObject* mod = PyImport_ImportModule("orbit.core.spacecharge"); - PyObject* pyUniformEllipsoidFieldCalculator = PyObject_CallMethod(mod,const_cast("UniformEllipsoidFieldCalculator"),const_cast("")); - //delete the c++ reference to the internal UniformEllipsoidFieldCalculator inside pyUniformEllipsoidFieldCalculator and assign the new one - delete ((UniformEllipsoidFieldCalculator*)((pyORBIT_Object*) pyUniformEllipsoidFieldCalculator)->cpp_obj); - ((pyORBIT_Object*) pyUniformEllipsoidFieldCalculator)->cpp_obj = cpp_ellipseFieldCalc; - cpp_ellipseFieldCalc->setPyWrapper(pyUniformEllipsoidFieldCalculator); - Py_INCREF(cpp_ellipseFieldCalc->getPyWrapper()); - Py_DECREF(mod); - return pyUniformEllipsoidFieldCalculator; - } - - //getNEllipses() - returns the number of ellipses inside the Space Charge calculator - static PyObject* SpaceChargeCalcUnifEllipse_getNEllipses(PyObject *self, PyObject *args){ - pyORBIT_Object* pySpaceChargeCalcUnifEllipse = (pyORBIT_Object*) self; - SpaceChargeCalcUnifEllipse* cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse*) pySpaceChargeCalcUnifEllipse->cpp_obj; - return Py_BuildValue("i",cpp_SpaceChargeCalcUnifEllipse->getNEllipses()); - } - - //calculateField(x,y,z) - calculates the fileds from all ellipses - static PyObject* SpaceChargeCalcUnifEllipse_calculateField(PyObject *self, PyObject *args){ - pyORBIT_Object* pySpaceChargeCalcUnifEllipse = (pyORBIT_Object*) self; - SpaceChargeCalcUnifEllipse* cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse*) pySpaceChargeCalcUnifEllipse->cpp_obj; - double x,y,z,ex,ey,ez; - if(!PyArg_ParseTuple(args,"ddd:calculateField",&x,&y,&z)){ - ORBIT_MPI_Finalize("PySpaceChargeCalcUnifEllipse.calculateField(x,y,z) - method needs parameters."); - } - cpp_SpaceChargeCalcUnifEllipse->calculateField(x,y,z,ex,ey,ez); - return Py_BuildValue("(ddd)",ex,ey,ez);; - } - - //----------------------------------------------------- - //destructor for python SpaceChargeCalcUnifEllipse class (__del__ method). - //----------------------------------------------------- - static void SpaceChargeCalcUnifEllipse_del(pyORBIT_Object* self){ - SpaceChargeCalcUnifEllipse* cpp_SpaceChargeCalcUnifEllipse = (SpaceChargeCalcUnifEllipse*) self->cpp_obj; - if(cpp_SpaceChargeCalcUnifEllipse != NULL){ - delete cpp_SpaceChargeCalcUnifEllipse; - } - self->ob_base.ob_type->tp_free((PyObject*)self); - } - - // defenition of the methods of the python SpaceChargeCalcUnifEllipse wrapper class - // they will be vailable from python level - static PyMethodDef SpaceChargeCalcUnifEllipseClassMethods[] = { - { "trackBunch", SpaceChargeCalcUnifEllipse_trackBunch, METH_VARARGS,"track the bunch - trackBunch(pyBunch,length)"}, - { "getNEllipses", SpaceChargeCalcUnifEllipse_getNEllipses, METH_VARARGS,"returns the number of ellipses inside the Space Charge calculator"}, - { "getEllipsFieldCalculator", SpaceChargeCalcUnifEllipse_getEllipsFieldCalculator, METH_VARARGS,"returns the ellipse field object with particular index"}, - { "calculateField", SpaceChargeCalcUnifEllipse_calculateField, METH_VARARGS,"calculates the fileds ex,ey,ez from all ellipses"}, - {NULL} - }; - - // defenition of the memebers of the python SpaceChargeCalcUnifEllipse wrapper class - // they will be vailable from python level - static PyMemberDef SpaceChargeCalcUnifEllipseClassMembers [] = { - {NULL} - }; - - //new python SpaceChargeCalcUnifEllipse wrapper type definition - static PyTypeObject pyORBIT_SpaceChargeCalcUnifEllipse_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "SpaceChargeCalcUnifEllipse", /*tp_name*/ - sizeof(pyORBIT_Object), /*tp_basicsize*/ - 0, /*tp_itemsize*/ - (destructor) SpaceChargeCalcUnifEllipse_del , /*tp_dealloc*/ - 0, /*tp_print*/ - 0, /*tp_getattr*/ - 0, /*tp_setattr*/ - 0, /*tp_compare*/ - 0, /*tp_repr*/ - 0, /*tp_as_number*/ - 0, /*tp_as_sequence*/ - 0, /*tp_as_mapping*/ - 0, /*tp_hash */ - 0, /*tp_call*/ - 0, /*tp_str*/ - 0, /*tp_getattro*/ - 0, /*tp_setattro*/ - 0, /*tp_as_buffer*/ - Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ - "The SpaceChargeCalcUnifEllipse python wrapper", /* tp_doc */ - 0, /* tp_traverse */ - 0, /* tp_clear */ - 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ - 0, /* tp_iter */ - 0, /* tp_iternext */ - SpaceChargeCalcUnifEllipseClassMethods, /* tp_methods */ - SpaceChargeCalcUnifEllipseClassMembers, /* tp_members */ - 0, /* tp_getset */ - 0, /* tp_base */ - 0, /* tp_dict */ - 0, /* tp_descr_get */ - 0, /* tp_descr_set */ - 0, /* tp_dictoffset */ - (initproc) SpaceChargeCalcUnifEllipse_init, /* tp_init */ - 0, /* tp_alloc */ - SpaceChargeCalcUnifEllipse_new, /* tp_new */ - }; - - //-------------------------------------------------- - //Initialization function of the pySpaceChargeCalcUnifEllipse class - //It will be called from SpaceCharge wrapper initialization - //-------------------------------------------------- - void initSpaceChargeCalcUniformEllipse(PyObject* module){ - if (PyType_Ready(&pyORBIT_SpaceChargeCalcUnifEllipse_Type) < 0) return; - Py_INCREF(&pyORBIT_SpaceChargeCalcUnifEllipse_Type); - PyModule_AddObject(module, "SpaceChargeCalcUnifEllipse", (PyObject *)&pyORBIT_SpaceChargeCalcUnifEllipse_Type); - } - -#ifdef __cplusplus -} -#endif - -//end of namespace wrap_spacecharge -} diff --git a/src/spacecharge/wrap_uniform_ellipsoid_field_calculator.cc b/src/spacecharge/wrap_uniform_ellipsoid_field_calculator.cc index 37e0b3d1..dbe12ca6 100644 --- a/src/spacecharge/wrap_uniform_ellipsoid_field_calculator.cc +++ b/src/spacecharge/wrap_uniform_ellipsoid_field_calculator.cc @@ -3,149 +3,148 @@ #include "UniformEllipsoidFieldCalculator.hh" -#include "wrap_uniform_ellipsoid_field_calculator.hh" #include "wrap_spacecharge.hh" +#include "wrap_uniform_ellipsoid_field_calculator.hh" #include using namespace OrbitUtils; -namespace wrap_spacecharge{ +namespace wrap_spacecharge { #ifdef __cplusplus extern "C" { #endif - //--------------------------------------------------------- - //Python UniformEllipsoidFieldCalculator class definition - //--------------------------------------------------------- - - //constructor for python class wrapping UniformEllipsoidFieldCalculator instance - //It never will be called directly - static PyObject* UniformEllipsoidFieldCalculator_new(PyTypeObject *type, PyObject *args, PyObject *kwds) - { - pyORBIT_Object* self; - self = (pyORBIT_Object *) type->tp_alloc(type, 0); - self->cpp_obj = NULL; - return (PyObject *) self; - } - - //initializator for python UniformEllipsoidFieldCalculator class - //this is implementation of the __init__ method - static int UniformEllipsoidFieldCalculator_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ - self->cpp_obj = new UniformEllipsoidFieldCalculator(); - ((UniformEllipsoidFieldCalculator*) self->cpp_obj)->setPyWrapper((PyObject*) self); - return 0; - } - - /** Sets the half-axis of the ellipsoid and maximal values of radius*/ - static PyObject* UniformEllipsoidFieldCalculator_setEllipsoid(PyObject *self, PyObject *args){ - pyORBIT_Object* pyUniformEllipsoidFieldCalculator = (pyORBIT_Object*) self; - UniformEllipsoidFieldCalculator* cpp_UniformEllipsoidFieldCalculator = (UniformEllipsoidFieldCalculator*) pyUniformEllipsoidFieldCalculator->cpp_obj; - double a,b,c,r_max; - if(!PyArg_ParseTuple(args,"dddd:setEllipsoid",&a,&b,&c,&r_max)){ - ORBIT_MPI_Finalize("PyUniformEllipsoidFieldCalculator.setEllipsoid(a,b,c,r_max) - method needs parameters."); - } - cpp_UniformEllipsoidFieldCalculator->setEllipsoid(a,b,c,r_max); - Py_INCREF(Py_None); +//--------------------------------------------------------- +// Python UniformEllipsoidFieldCalculator class definition +//--------------------------------------------------------- + +// constructor for python class wrapping UniformEllipsoidFieldCalculator instance +// It never will be called directly +static PyObject *UniformEllipsoidFieldCalculator_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { + pyORBIT_Object *self; + self = (pyORBIT_Object *)type->tp_alloc(type, 0); + self->cpp_obj = NULL; + return (PyObject *)self; +} + +// initializator for python UniformEllipsoidFieldCalculator class +// this is implementation of the __init__ method +static int UniformEllipsoidFieldCalculator_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds) { + self->cpp_obj = new UniformEllipsoidFieldCalculator(); + ((UniformEllipsoidFieldCalculator *)self->cpp_obj)->setPyWrapper((PyObject *)self); + return 0; +} + +/** Sets the half-axis of the ellipsoid and maximal values of radius*/ +static PyObject *UniformEllipsoidFieldCalculator_setEllipsoid(PyObject *self, PyObject *args) { + pyORBIT_Object *pyUniformEllipsoidFieldCalculator = (pyORBIT_Object *)self; + UniformEllipsoidFieldCalculator *cpp_UniformEllipsoidFieldCalculator = (UniformEllipsoidFieldCalculator *)pyUniformEllipsoidFieldCalculator->cpp_obj; + double a, b, c, r_max; + if (!PyArg_ParseTuple(args, "dddd:setEllipsoid", &a, &b, &c, &r_max)) { + ORBIT_MPI_Finalize("PyUniformEllipsoidFieldCalculator.setEllipsoid(a,b,c,r_max) - method needs parameters."); + } + cpp_UniformEllipsoidFieldCalculator->setEllipsoid(a, b, c, r_max); + Py_INCREF(Py_None); return Py_None; - } - - /** Calculates the field components */ - static PyObject* UniformEllipsoidFieldCalculator_calcField(PyObject *self, PyObject *args){ - pyORBIT_Object* pyUniformEllipsoidFieldCalculator = (pyORBIT_Object*) self; - UniformEllipsoidFieldCalculator* cpp_UniformEllipsoidFieldCalculator = (UniformEllipsoidFieldCalculator*) pyUniformEllipsoidFieldCalculator->cpp_obj; - double x,y,z; - if(!PyArg_ParseTuple(args,"ddd:calcField",&x,&y,&z)){ - ORBIT_MPI_Finalize("PyUniformEllipsoidFieldCalculator.calcField(x,y,z) - method needs parameters."); - } - double x2,y2,z2; - double ex,ey,ez; - x2 = x*x; y2 = y*y; z2 = z*z; - cpp_UniformEllipsoidFieldCalculator->calcField(x,y,z,x2,y2,z2,ex,ey,ez); - return Py_BuildValue("(ddd)",ex,ey,ez); - } - - //----------------------------------------------------- - //destructor for python UniformEllipsoidFieldCalculator class (__del__ method). - //----------------------------------------------------- - static void UniformEllipsoidFieldCalculator_del(pyORBIT_Object* self){ - UniformEllipsoidFieldCalculator* cpp_UniformEllipsoidFieldCalculator = (UniformEllipsoidFieldCalculator*) self->cpp_obj; - if(cpp_UniformEllipsoidFieldCalculator != NULL){ - delete cpp_UniformEllipsoidFieldCalculator; - } - self->ob_base.ob_type->tp_free((PyObject*)self); - } - - // defenition of the methods of the python UniformEllipsoidFieldCalculator wrapper class - // they will be vailable from python level - static PyMethodDef UniformEllipsoidFieldCalculatorClassMethods[] = { - { "setEllipsoid", UniformEllipsoidFieldCalculator_setEllipsoid, METH_VARARGS,"sets the half-axis of the ellipsoid and maximal values of radius"}, - { "calcField", UniformEllipsoidFieldCalculator_calcField, METH_VARARGS,"returns (ex,ey,ez) for (x,y,z) input"}, - {NULL} - }; - - // defenition of the memebers of the python UniformEllipsoidFieldCalculator wrapper class - // they will be vailable from python level - static PyMemberDef UniformEllipsoidFieldCalculatorClassMembers [] = { - {NULL} - }; - - //new python UniformEllipsoidFieldCalculator wrapper type definition - static PyTypeObject pyORBIT_UniformEllipsoidFieldCalculator_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "UniformEllipsoidFieldCalculator", /*tp_name*/ - sizeof(pyORBIT_Object), /*tp_basicsize*/ - 0, /*tp_itemsize*/ - (destructor) UniformEllipsoidFieldCalculator_del , /*tp_dealloc*/ - 0, /*tp_print*/ - 0, /*tp_getattr*/ - 0, /*tp_setattr*/ - 0, /*tp_compare*/ - 0, /*tp_repr*/ - 0, /*tp_as_number*/ - 0, /*tp_as_sequence*/ - 0, /*tp_as_mapping*/ - 0, /*tp_hash */ - 0, /*tp_call*/ - 0, /*tp_str*/ - 0, /*tp_getattro*/ - 0, /*tp_setattro*/ - 0, /*tp_as_buffer*/ - Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ - "The UniformEllipsoidFieldCalculator python wrapper", /* tp_doc */ - 0, /* tp_traverse */ - 0, /* tp_clear */ - 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ - 0, /* tp_iter */ - 0, /* tp_iternext */ - UniformEllipsoidFieldCalculatorClassMethods, /* tp_methods */ - UniformEllipsoidFieldCalculatorClassMembers, /* tp_members */ - 0, /* tp_getset */ - 0, /* tp_base */ - 0, /* tp_dict */ - 0, /* tp_descr_get */ - 0, /* tp_descr_set */ - 0, /* tp_dictoffset */ - (initproc) UniformEllipsoidFieldCalculator_init, /* tp_init */ - 0, /* tp_alloc */ - UniformEllipsoidFieldCalculator_new, /* tp_new */ - }; - - //-------------------------------------------------- - //Initialization function of the pyUniformEllipsoidFieldCalculator class - //It will be called from SpaceCharge wrapper initialization - //-------------------------------------------------- - void initUniformEllipsoidFieldCalculator(PyObject* module){ - if (PyType_Ready(&pyORBIT_UniformEllipsoidFieldCalculator_Type) < 0) return; - Py_INCREF(&pyORBIT_UniformEllipsoidFieldCalculator_Type); - PyModule_AddObject(module, "UniformEllipsoidFieldCalculator", (PyObject *)&pyORBIT_UniformEllipsoidFieldCalculator_Type); - } +} + +/** Calculates the field components */ +static PyObject *UniformEllipsoidFieldCalculator_calcField(PyObject *self, PyObject *args) { + pyORBIT_Object *pyUniformEllipsoidFieldCalculator = (pyORBIT_Object *)self; + UniformEllipsoidFieldCalculator *cpp_UniformEllipsoidFieldCalculator = (UniformEllipsoidFieldCalculator *)pyUniformEllipsoidFieldCalculator->cpp_obj; + double x, y, z; + if (!PyArg_ParseTuple(args, "ddd:calcField", &x, &y, &z)) { + ORBIT_MPI_Finalize("PyUniformEllipsoidFieldCalculator.calcField(x,y,z) - method needs parameters."); + } + double x2, y2, z2; + double ex, ey, ez; + x2 = x * x; + y2 = y * y; + z2 = z * z; + cpp_UniformEllipsoidFieldCalculator->calcField(x, y, z, x2, y2, z2, ex, ey, ez); + return Py_BuildValue("(ddd)", ex, ey, ez); +} + +//----------------------------------------------------- +// destructor for python UniformEllipsoidFieldCalculator class (__del__ method). +//----------------------------------------------------- +static void UniformEllipsoidFieldCalculator_del(pyORBIT_Object *self) { + UniformEllipsoidFieldCalculator *cpp_UniformEllipsoidFieldCalculator = (UniformEllipsoidFieldCalculator *)self->cpp_obj; + if (cpp_UniformEllipsoidFieldCalculator != NULL) { + delete cpp_UniformEllipsoidFieldCalculator; + } + self->ob_base.ob_type->tp_free((PyObject *)self); +} + +// defenition of the methods of the python UniformEllipsoidFieldCalculator wrapper class +// they will be vailable from python level +static PyMethodDef UniformEllipsoidFieldCalculatorClassMethods[] = { + {"setEllipsoid", UniformEllipsoidFieldCalculator_setEllipsoid, METH_VARARGS, "sets the half-axis of the ellipsoid and maximal values of radius"}, + {"calcField", UniformEllipsoidFieldCalculator_calcField, METH_VARARGS, "returns (ex,ey,ez) for (x,y,z) input"}, + {NULL}}; + +// defenition of the memebers of the python UniformEllipsoidFieldCalculator wrapper class +// they will be vailable from python level +static PyMemberDef UniformEllipsoidFieldCalculatorClassMembers[] = { + {NULL}}; + +// new python UniformEllipsoidFieldCalculator wrapper type definition +static PyTypeObject pyORBIT_UniformEllipsoidFieldCalculator_Type = { + PyVarObject_HEAD_INIT(NULL, 0) "UniformEllipsoidFieldCalculator", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor)UniformEllipsoidFieldCalculator_del, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ + "The UniformEllipsoidFieldCalculator python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + UniformEllipsoidFieldCalculatorClassMethods, /* tp_methods */ + UniformEllipsoidFieldCalculatorClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)UniformEllipsoidFieldCalculator_init, /* tp_init */ + 0, /* tp_alloc */ + UniformEllipsoidFieldCalculator_new, /* tp_new */ +}; + +//-------------------------------------------------- +// Initialization function of the pyUniformEllipsoidFieldCalculator class +// It will be called from SpaceCharge wrapper initialization +//-------------------------------------------------- +void initUniformEllipsoidFieldCalculator(PyObject *module) { + if (PyType_Ready(&pyORBIT_UniformEllipsoidFieldCalculator_Type) < 0) + return; + Py_INCREF(&pyORBIT_UniformEllipsoidFieldCalculator_Type); + PyModule_AddObject(module, "UniformEllipsoidFieldCalculator", (PyObject *)&pyORBIT_UniformEllipsoidFieldCalculator_Type); +} #ifdef __cplusplus } #endif -//end of namespace wrap_spacecharge -} +// end of namespace wrap_spacecharge +} // namespace wrap_spacecharge diff --git a/tests/py/orbit/test_sc_unif_ellipse.py b/tests/py/orbit/test_sc_unif_ellipse.py new file mode 100644 index 00000000..99d5d4d7 --- /dev/null +++ b/tests/py/orbit/test_sc_unif_ellipse.py @@ -0,0 +1,77 @@ +import math + +import pytest + +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse2D +from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse + + +def test_uniform_ellipsoid_calculator_names(): + assert SpaceChargeCalcUnifEllipse2D().getNEllipses() == 1 + assert SpaceChargeCalcUnifEllipse().getNEllipses() == 1 + + +def test_uniform_ellipse_2d_kick_matches_envelope_formula(): + bunch = Bunch() + bunch.mass(0.93827231) + bunch.charge(1.0) + bunch.getSyncParticle().kinEnergy(1.0) + + radius_x = 0.004 + radius_y = 0.002 + bunch_length = 0.12 + z_rms = bunch_length / math.sqrt(12.0) + phi = 0.37 + cos_phi = math.cos(phi) + sin_phi = math.sin(phi) + local_coordinates = ( + (+radius_x / math.sqrt(2.0), 0.0, +z_rms), + (-radius_x / math.sqrt(2.0), 0.0, -z_rms), + (0.0, +radius_y / math.sqrt(2.0), +z_rms), + (0.0, -radius_y / math.sqrt(2.0), -z_rms), + ) + coordinates = [ + (cos_phi * x + sin_phi * y, -sin_phi * x + cos_phi * y, z) + for x, y, z in local_coordinates + ] + for x, y, z in coordinates: + bunch.addParticle(x, 0.0, y, 0.0, z, 0.0) + + macro_size = 2.5e8 + bunch.macroSize(macro_size) + intensity = macro_size * len(coordinates) + length = 0.25 + + beta = bunch.getSyncParticle().beta() + gamma = bunch.getSyncParticle().gamma() + perveance = ( + 2.0 + * intensity + * bunch.classicalRadius() + / (beta**2 * gamma**3 * bunch_length) + ) + kappa_factor = 2.0 * perveance / (radius_x + radius_y) + + calculator = SpaceChargeCalcUnifEllipse2D() + calculator.trackBunch(bunch, length) + + for index, (x, y, _) in enumerate(local_coordinates): + xp = kappa_factor * length * x / radius_x + yp = kappa_factor * length * y / radius_y + expected_xp = cos_phi * xp + sin_phi * yp + expected_yp = -sin_phi * xp + cos_phi * yp + assert bunch.xp(index) == pytest.approx(expected_xp) + assert bunch.yp(index) == pytest.approx(expected_yp) + assert all(bunch.dE(index) == 0.0 for index in range(bunch.getSize())) + + x = 2.0 * radius_x + x_lab = cos_phi * x + y_lab = -sin_phi * x + ex, ey, ez = calculator.calculateField(x_lab, y_lab, 0.0) + lambda_value = x**2 - radius_x**2 + radius_y_lambda = math.sqrt(radius_y**2 + lambda_value) + ex_local = 2.0 * (intensity / bunch_length) / (x + radius_y_lambda) + assert ex == pytest.approx(cos_phi * ex_local) + assert ey == pytest.approx(-sin_phi * ex_local) + assert ez == 0.0 diff --git a/tests/py/orbit/test_sc_unif_ellipse_nodes.py b/tests/py/orbit/test_sc_unif_ellipse_nodes.py new file mode 100644 index 00000000..09fe185a --- /dev/null +++ b/tests/py/orbit/test_sc_unif_ellipse_nodes.py @@ -0,0 +1,32 @@ +import math +import random + +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse +from orbit.space_charge.sc2p5d import SCUnifEllipse2D_AccNode +from orbit.space_charge.sc2p5d import setSCUnifEllipse2DAccNodes +from orbit.space_charge.sc3d import SCUnifEllipse_AccNode +from orbit.space_charge.sc3d import setSCUnifEllipseAccNodes + + +def test_unif_ellipse_node_track_bunch(): + bunch = Bunch() + bunch.mass(0.938) + bunch.getSyncParticle().kinEnergy(0.001) + + for i in range(10): + x = random.gauss(0.0, 0.001) + y = random.gauss(0.0, 0.001) + z = random.gauss(0.0, 0.001) + bunch.addParticle(x, 0.0, y, 0.0, z, 0.0) + + params_dict = {"bunch": bunch} + + calculator = SpaceChargeCalcUnifEllipse() + node = SCUnifEllipse2D_AccNode(calculator) + node.track(params_dict) + + +def test_unif_ellipse_node_setters_are_exported(): + assert callable(setSCUnifEllipseAccNodes) + assert callable(setSCUnifEllipse2DAccNodes) From 1a40c15ef47c715c6e4601bfe054437724a08180 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Sat, 5 Sep 2026 00:59:18 -0400 Subject: [PATCH 2/6] Add SCUnifEllipse benchmark (FODO lattice) --- examples/SpaceCharge/plot.py | 133 +++++++++ examples/SpaceCharge/style.mplstyle | 8 + .../SpaceCharge/test_sc_unif_ellipse_fodo.py | 276 ++++++++++++++++++ examples/SpaceCharge/utils.py | 60 ++++ 4 files changed, 477 insertions(+) create mode 100644 examples/SpaceCharge/plot.py create mode 100644 examples/SpaceCharge/style.mplstyle create mode 100644 examples/SpaceCharge/test_sc_unif_ellipse_fodo.py create mode 100644 examples/SpaceCharge/utils.py diff --git a/examples/SpaceCharge/plot.py b/examples/SpaceCharge/plot.py new file mode 100644 index 00000000..a6c5c1aa --- /dev/null +++ b/examples/SpaceCharge/plot.py @@ -0,0 +1,133 @@ +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as patches + + +def calc_rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]: + """Return rms ellipse dimensions and orientation.""" + i, j = (0, 1) + + sii = cov_matrix[i, i] + sjj = cov_matrix[j, j] + sij = cov_matrix[i, j] + + angle = -0.5 * np.arctan2(2.0 * sij, sii - sjj) + + _sin = np.sin(angle) + _cos = np.cos(angle) + _sin2 = _sin**2 + _cos2 = _cos**2 + + c1 = np.sqrt(abs(sii * _cos2 + sjj * _sin2 - 2 * sij * _sin * _cos)) + c2 = np.sqrt(abs(sii * _sin2 + sjj * _cos2 + 2 * sij * _sin * _cos)) + + return (c1, c2, angle) + + +def plot_ellipse( + r1: float = 1.0, + r2: float = 1.0, + angle: float = 0.0, + center: tuple[float, float] = None, + ax=None, + **kws, +): + kws.setdefault("fill", False) + kws.setdefault("color", "black") + kws.setdefault("lw", 1.25) + + if center is None: + center = (0.0, 0.0) + + d1 = r1 * 2.0 + d2 = r2 * 2.0 + angle = -np.degrees(angle) + + ax.add_patch(patches.Ellipse(center, d1, d2, angle=angle, **kws)) + return ax + + +def plot_rms_ellipse( + cov_matrix: np.ndarray, + level: float = 1.0, + ax=None, + **ellipse_kws, +): + """Plot rms ellipse from 2 x 2 covariance matrix.""" + r1, r2, angle = calc_rms_ellipse_params(cov_matrix) + plot_ellipse(r1 * level, r2 * level, angle=angle, ax=ax, **ellipse_kws) + return ax + + +def plot_corner( + particles: np.ndarray, + limits: list[tuple[float, float]] = None, + bins: int = 64, + labels: list[str] = None, + blur: float = None, +) -> tuple: + """Generate corner plot.""" + ndim = particles.shape[1] + + if limits is None: + xmax = np.max(particles, axis=0) + xmin = np.min(particles, axis=0) + limits = list(zip(xmin, xmax)) + + if labels is None: + labels = ndim * [""] + + fig, axs = plt.subplots( + ncols=ndim, nrows=ndim, sharex=None, sharey=None, figsize=(7, 7) + ) + for i in range(ndim): + for j in range(ndim): + axis = (j, i) + ax = axs[i, j] + if i > j: + values, edges = np.histogramdd( + particles[:, axis], bins=bins, range=[limits[k] for k in axis] + ) + if blur: + values = scipy.ndimage.gaussian_filter(values, sigma=blur) + ax.pcolormesh( + edges[0], + edges[1], + values.T, + linewidth=0.0, + rasterized=True, + shading="auto", + ) + elif i == j: + values, edges = np.histogram( + particles[:, i], bins=bins, range=limits[i] + ) + if blur: + values = scipy.ndimage.gaussian_filter(values, sigma=blur) + ax.stairs(values, edges, lw=1.5, color="black") + else: + ax.axis("off") + + for i in range(0, ndim - 1): + for j in range(0, ndim): + axs[i, j].set_xticklabels([]) + for i in range(0, ndim): + for j in range(1, ndim): + axs[i, j].set_yticklabels([]) + + for ax in axs.flat: + for loc in ["top", "right"]: + ax.spines[loc].set_visible(False) + + for i, label in enumerate(labels): + axs[-1, i].set_xlabel(label) + for i, label in enumerate(labels[1:], start=1): + axs[i, 0].set_ylabel(label) + + axs[0, 0].set_yticklabels([]) + axs[0, 0].set_ylabel(None) + + fig.align_ylabels() + fig.align_xlabels() + + return fig, axs diff --git a/examples/SpaceCharge/style.mplstyle b/examples/SpaceCharge/style.mplstyle new file mode 100644 index 00000000..71f36605 --- /dev/null +++ b/examples/SpaceCharge/style.mplstyle @@ -0,0 +1,8 @@ +axes.linewidth: 1.25 +axes.titlesize: "medium" +image.cmap: "Greys" +figure.constrained_layout.use: True +savefig.dpi: 300 +savefig.format: "png" +xtick.minor.visible: True +ytick.minor.visible: True diff --git a/examples/SpaceCharge/test_sc_unif_ellipse_fodo.py b/examples/SpaceCharge/test_sc_unif_ellipse_fodo.py new file mode 100644 index 00000000..79d0e495 --- /dev/null +++ b/examples/SpaceCharge/test_sc_unif_ellipse_fodo.py @@ -0,0 +1,276 @@ +"""Test 2D envelope tracker in FODO lattice.""" + +import argparse +import copy +import math +import os +import pathlib +import time + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.spacecharge import SpaceChargeCalc2p5D +from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse2D +from orbit.bunch_generators import TwissContainer +from orbit.bunch_generators import KVDist2D +from orbit.bunch_generators import WaterBagDist2D +from orbit.bunch_generators import GaussDist2D +from orbit.bunch_utils import collect_bunch +from orbit.lattice import AccNode +from orbit.lattice import AccLattice +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.space_charge.sc2p5d import setSCUnifEllipse2DAccNodes +from orbit.teapot import DriftTEAPOT +from orbit.teapot import QuadTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import build_rotation_matrix_xy +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--bunch-length", type=float, default=5.0) + parser.add_argument("--kin-energy", type=float, default=0.0025) + parser.add_argument("--intensity", type=float, default=3e9) + + parser.add_argument("--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"]) + parser.add_argument("--mismatch-x", type=float, default=0.0) + parser.add_argument("--mismatch-y", type=float, default=0.0) + parser.add_argument("--tilt", type=float, default=0) + + parser.add_argument("--nslice", type=int, default=10) + parser.add_argument("--kq", type=float, default=0.25) + + parser.add_argument("--nparts", type=int, default=100_000) + parser.add_argument("--nturns", type=int, default=25) + parser.add_argument("--sc", type=int, default=0) + parser.add_argument("--sc-grid-res", type=int, default=128) + parser.add_argument("--sc-ellipse-n", type=int, default=1) + return parser.parse_args() + + +def make_lattice(args: argparse.Namespace) -> AccLattice: + nodes = [ + QuadTEAPOT(length=0.5, kq=+args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=1.0, kq=-args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=0.5, kq=+args.kq), + ] + + lattice = TEAPOT_Lattice() + for node in nodes: + node.setnParts(args.nslice) + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + lattice.addNode(node) + lattice.initialize() + return lattice + + +def make_empy_bunch(args: argparse.Namespace) -> Bunch: + bunch = Bunch() + bunch.mass(mass_proton) + bunch.getSyncParticle().kinEnergy(args.kin_energy) + return bunch + + +def make_bunch(args: argparse.Namespace) -> Bunch: + bunch = make_empy_bunch(args) + lattice = make_lattice(args) + + matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) + matrix_lattice_params = matrix_lattice.getRingParametersDict() + + alpha_x = matrix_lattice_params["alpha x"] + alpha_y = matrix_lattice_params["alpha y"] + beta_x = matrix_lattice_params["beta x [m]"] + beta_y = matrix_lattice_params["beta y [m]"] + eps_x = 0.25e-06 + eps_y = eps_x + + twiss_x = TwissContainer(alpha_x, beta_x, eps_x) + twiss_y = TwissContainer(alpha_y, beta_y, eps_y) + + if args.dist == "kv": + dist = KVDist2D(twiss_x, twiss_y) + elif args.dist == "waterbag": + dist = WaterBagDist2D(twiss_x, twiss_y) + elif args.dist == "gauss": + dist = GaussDist2D(twiss_x, twiss_y) + else: + raise ValueError + + particles = np.zeros((args.nparts, 6)) + for i in range(args.nparts): + particles[i, :4] = dist.getCoordinates() + particles[i, 4] = args.bunch_length * np.random.uniform(-1.0, 1.0) + + matrix = build_rotation_matrix_xy(math.radians(args.tilt)) + particles[:, :4] = particles[:, :4] @ matrix.T + + particles[:, 0] *= (1.0 + args.mismatch_x) + particles[:, 2] *= (1.0 + args.mismatch_y) + + for i in range(args.nparts): + bunch.addParticle(*particles[i]) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + return bunch + + +def get_bunch_cov(bunch: Bunch) -> np.ndarray: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + return cov_matrix + + +def track(lattice: TEAPOT_Lattice, bunch: Bunch, nturns: int) -> dict: + bunch_out = Bunch() + bunch.copyBunchTo(bunch_out) + + history = {"xrms": [], "yrms": []} + for turn in range(nturns): + if turn > 0: + lattice.trackBunch(bunch_out) + + cov_matrix = get_bunch_cov(bunch_out) + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f}") + + particles_out = collect_bunch(bunch_out)["coords"] + particles_out = particles_out[:, (0, 1, 2, 3)] + particles_out *= 1000.0 + return { + "particles": particles_out, + "history": history, + } + + +def main(args: argparse.Namespace) -> None: + path = pathlib.Path(__file__) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) + os.makedirs(output_dir, exist_ok=True) + + results = {} + for key in ["SC2p5D", "SCUnifEllipse"]: + results[key] = {} + + # Make initial bunch + bunch = make_bunch(args) + + # Track bunch with SC2p5D space charge nodes + lattice = make_lattice(args) + if args.sc: + sc_calc = SpaceChargeCalc2p5D(args.sc_grid_res, args.sc_grid_res, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + print("TRACK SC2p5D") + results["SC2p5D"] = track(lattice, bunch, nturns=args.nturns) + + # Track bunch with SCUnifEllipse space charge nodes + lattice = make_lattice(args) + if args.sc: + sc_calc = SpaceChargeCalcUnifEllipse2D(args.sc_ellipse_n) + sc_nodes = setSCUnifEllipse2DAccNodes(lattice, sc_path_length_min, sc_calc) + + print("TRACK SCUnifEllipse2D") + results["SCUnifEllipse"] = track(lattice, bunch, nturns=args.nturns) + + + # Analysis + # ------------------------------------------------------------------------------ + + for model, result in results.items(): + history = result["history"] + for key in history: + history[key] = np.array(history[key]) + + # Print errors + for key in results["SCUnifEllipse"]["history"]: + deltas = results["SCUnifEllipse"]["history"][key] - results["SC2p5D"]["history"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + + # Plot rms bunch sizes + for key in ["xrms", "yrms"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["SC2p5D", "SCUnifEllipse"]): + plot_kws = {} + plot_kws["color"] = ["black", "red"][i] + plot_kws["lw"] = [None, 0][i] + ax.plot(results[model]["history"][key], marker=".", label=model, **plot_kws) + ax.set_ylim(0.0, ax.get_ylim()[1] * 2.0) + ax.set_xlabel("Turn") + ax.set_ylabel("RMS [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + + # Set plot limits + particles = results["SC2p5D"]["particles"] + xmax = 4.0 * np.std(particles, axis=0) + limits = list(zip(-xmax, xmax)) + labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]"] + + # Plot x-x' + fig, axs = plt.subplots(figsize=(6, 3), ncols=2, sharex=True, sharey=True) + for ax, model in zip(axs, results): + particles = results[model]["particles"] + ax.hist2d(particles[:, 0], particles[:, 1], bins=64, range=[limits[0], limits[1]]) + ax.set_xlabel(labels[0]) + ax.set_ylabel(labels[1]) + ax.set_title(model) + plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) + plt.close() + + # Plot y-y' + fig, axs = plt.subplots(figsize=(6, 3), ncols=2, sharex=True, sharey=True) + for ax, model in zip(axs, results): + particles = results[model]["particles"] + ax.hist2d(particles[:, 2], particles[:, 3], bins=64, range=[limits[2], limits[3]]) + ax.set_xlabel(labels[2]) + ax.set_ylabel(labels[3]) + ax.set_title(model) + plt.savefig(os.path.join(output_dir, "fig_dist_y_yp")) + plt.close() + + # Plot corner + for model in results: + particles = results[model]["particles"] + fig, axs = plot_corner( + particles, + limits=limits, + bins=64, + labels=labels, + ) + plt.savefig(os.path.join(output_dir, f"fig_dist_corner_{model}")) + plt.close() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/SpaceCharge/utils.py b/examples/SpaceCharge/utils.py new file mode 100644 index 00000000..daf2a6c7 --- /dev/null +++ b/examples/SpaceCharge/utils.py @@ -0,0 +1,60 @@ +import numpy as np + + +def gen_dist_gauss(size: int, dim: np.ndarray) -> np.ndarray: + return np.random.multivariate_normal( + mean=np.zeros(dim), + cov=np.eye(dim), + size=size, + ) + + +def gen_dist_kv(size: int, dim: int) -> np.ndarray: + X = np.random.normal(size=(size, dim)) + X /= np.linalg.norm(X, axis=1)[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist_waterbag(size: int, dim: int) -> np.ndarray: + X = gen_dist_kv(size, dim) + dim = X.shape[1] + r = np.random.uniform(size=size) ** (1.0 / dim) + X *= r[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist(size: int, cov_matrix: np.ndarray, name: str) -> np.ndarray: + dim = cov_matrix.shape[0] + if name == "kv": + X = gen_dist_kv(size, dim) + elif name == "waterbag": + X = gen_dist_waterbag(size, dim) + elif name == "gauss": + X = gen_dist_gauss(size, dim) + else: + raise ValueError(f"Invalid distribution name: {name}") + + L = np.linalg.cholesky(cov_matrix) + return np.matmul(X, L.T) + + +def build_rotation_matrix_xy(angle: float) -> np.ndarray: + cs = np.cos(angle) + sn = np.sin(angle) + + matrix = np.identity(4) + matrix[0, 0] = matrix[1, 1] = +cs + matrix[0, 2] = matrix[1, 3] = +sn + matrix[2, 0] = matrix[3, 1] = -sn + matrix[2, 2] = matrix[3, 3] = +cs + return matrix + + +def project_cov_matrix(cov_matrix: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: + cov_matrix_proj = np.zeros((len(axis), len(axis))) + for i in range(len(axis)): + for j in range(len(axis)): + cov_matrix_proj[i, j] = cov_matrix[axis[i], axis[j]] + return cov_matrix_proj From 039183083493ee184dd4d5b2d996808890dba062 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Sat, 5 Sep 2026 01:01:22 -0400 Subject: [PATCH 3/6] Move test to separate folder --- examples/SpaceCharge/{ => scUnifEllipse}/plot.py | 0 examples/SpaceCharge/{ => scUnifEllipse}/style.mplstyle | 0 .../SpaceCharge/{ => scUnifEllipse}/test_sc_unif_ellipse_fodo.py | 0 examples/SpaceCharge/{ => scUnifEllipse}/utils.py | 0 4 files changed, 0 insertions(+), 0 deletions(-) rename examples/SpaceCharge/{ => scUnifEllipse}/plot.py (100%) rename examples/SpaceCharge/{ => scUnifEllipse}/style.mplstyle (100%) rename examples/SpaceCharge/{ => scUnifEllipse}/test_sc_unif_ellipse_fodo.py (100%) rename examples/SpaceCharge/{ => scUnifEllipse}/utils.py (100%) diff --git a/examples/SpaceCharge/plot.py b/examples/SpaceCharge/scUnifEllipse/plot.py similarity index 100% rename from examples/SpaceCharge/plot.py rename to examples/SpaceCharge/scUnifEllipse/plot.py diff --git a/examples/SpaceCharge/style.mplstyle b/examples/SpaceCharge/scUnifEllipse/style.mplstyle similarity index 100% rename from examples/SpaceCharge/style.mplstyle rename to examples/SpaceCharge/scUnifEllipse/style.mplstyle diff --git a/examples/SpaceCharge/test_sc_unif_ellipse_fodo.py b/examples/SpaceCharge/scUnifEllipse/test_sc_unif_ellipse_fodo.py similarity index 100% rename from examples/SpaceCharge/test_sc_unif_ellipse_fodo.py rename to examples/SpaceCharge/scUnifEllipse/test_sc_unif_ellipse_fodo.py diff --git a/examples/SpaceCharge/utils.py b/examples/SpaceCharge/scUnifEllipse/utils.py similarity index 100% rename from examples/SpaceCharge/utils.py rename to examples/SpaceCharge/scUnifEllipse/utils.py From 4017dc3e3c24059da71e2aba7af0edbdc7c9fb1a Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Sat, 5 Sep 2026 01:26:46 -0400 Subject: [PATCH 4/6] Update plotting --- examples/SpaceCharge/scUnifEllipse/plot.py | 17 ++-- .../SpaceCharge/scUnifEllipse/style.mplstyle | 2 +- ..._sc_unif_ellipse_fodo.py => track_fodo.py} | 85 +++++++++++-------- 3 files changed, 61 insertions(+), 43 deletions(-) rename examples/SpaceCharge/scUnifEllipse/{test_sc_unif_ellipse_fodo.py => track_fodo.py} (78%) diff --git a/examples/SpaceCharge/scUnifEllipse/plot.py b/examples/SpaceCharge/scUnifEllipse/plot.py index a6c5c1aa..ea0a12a3 100644 --- a/examples/SpaceCharge/scUnifEllipse/plot.py +++ b/examples/SpaceCharge/scUnifEllipse/plot.py @@ -1,6 +1,15 @@ import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as patches +import matplotlib.colors as mcolors + + +def trunc_cmap(cmap_name: str, vmin: float = 0.0, vmax: float = 1.0, n: int = 100): + cmap = plt.get_cmap(cmap_name) + new_cmap = mcolors.LinearSegmentedColormap.from_list( + f"trunc({cmap.name},{vmin:.2f},{vmax:.2f})", cmap(np.linspace(vmin, vmax, n)) + ) + return new_cmap def calc_rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]: @@ -64,7 +73,7 @@ def plot_corner( limits: list[tuple[float, float]] = None, bins: int = 64, labels: list[str] = None, - blur: float = None, + mask: bool = False, ) -> tuple: """Generate corner plot.""" ndim = particles.shape[1] @@ -88,8 +97,8 @@ def plot_corner( values, edges = np.histogramdd( particles[:, axis], bins=bins, range=[limits[k] for k in axis] ) - if blur: - values = scipy.ndimage.gaussian_filter(values, sigma=blur) + if mask: + values = np.ma.masked_less_equal(values, 0.0) ax.pcolormesh( edges[0], edges[1], @@ -102,8 +111,6 @@ def plot_corner( values, edges = np.histogram( particles[:, i], bins=bins, range=limits[i] ) - if blur: - values = scipy.ndimage.gaussian_filter(values, sigma=blur) ax.stairs(values, edges, lw=1.5, color="black") else: ax.axis("off") diff --git a/examples/SpaceCharge/scUnifEllipse/style.mplstyle b/examples/SpaceCharge/scUnifEllipse/style.mplstyle index 71f36605..b0d08dbb 100644 --- a/examples/SpaceCharge/scUnifEllipse/style.mplstyle +++ b/examples/SpaceCharge/scUnifEllipse/style.mplstyle @@ -1,6 +1,6 @@ axes.linewidth: 1.25 axes.titlesize: "medium" -image.cmap: "Greys" +image.cmap: "viridis" figure.constrained_layout.use: True savefig.dpi: 300 savefig.format: "png" diff --git a/examples/SpaceCharge/scUnifEllipse/test_sc_unif_ellipse_fodo.py b/examples/SpaceCharge/scUnifEllipse/track_fodo.py similarity index 78% rename from examples/SpaceCharge/scUnifEllipse/test_sc_unif_ellipse_fodo.py rename to examples/SpaceCharge/scUnifEllipse/track_fodo.py index 79d0e495..fac68557 100644 --- a/examples/SpaceCharge/scUnifEllipse/test_sc_unif_ellipse_fodo.py +++ b/examples/SpaceCharge/scUnifEllipse/track_fodo.py @@ -43,7 +43,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--kin-energy", type=float, default=0.0025) parser.add_argument("--intensity", type=float, default=3e9) - parser.add_argument("--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"]) + parser.add_argument( + "--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"] + ) parser.add_argument("--mismatch-x", type=float, default=0.0) parser.add_argument("--mismatch-y", type=float, default=0.0) parser.add_argument("--tilt", type=float, default=0) @@ -56,6 +58,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--sc", type=int, default=0) parser.add_argument("--sc-grid-res", type=int, default=128) parser.add_argument("--sc-ellipse-n", type=int, default=1) + + parser.add_argument("--plot-bins", type=int, default=64) + parser.add_argument("--plot-mask", action="store_true") return parser.parse_args() @@ -119,8 +124,8 @@ def make_bunch(args: argparse.Namespace) -> Bunch: matrix = build_rotation_matrix_xy(math.radians(args.tilt)) particles[:, :4] = particles[:, :4] @ matrix.T - particles[:, 0] *= (1.0 + args.mismatch_x) - particles[:, 2] *= (1.0 + args.mismatch_y) + particles[:, 0] *= 1.0 + args.mismatch_x + particles[:, 2] *= 1.0 + args.mismatch_y for i in range(args.nparts): bunch.addParticle(*particles[i]) @@ -146,19 +151,25 @@ def track(lattice: TEAPOT_Lattice, bunch: Bunch, nturns: int) -> dict: bunch_out = Bunch() bunch.copyBunchTo(bunch_out) - history = {"xrms": [], "yrms": []} + history = {"rms_x": [], "rms_y": [], "eps_x": [], "eps_y": []} for turn in range(nturns): if turn > 0: lattice.trackBunch(bunch_out) - cov_matrix = get_bunch_cov(bunch_out) - xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) - yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + cov_matrix = 1e6 * get_bunch_cov(bunch_out) + x_rms = math.sqrt(cov_matrix[0, 0]) + y_rms = math.sqrt(cov_matrix[2, 2]) + eps_x = np.sqrt(np.linalg.det(cov_matrix[0:2, 0:2])) + eps_y = np.sqrt(np.linalg.det(cov_matrix[2:4, 2:4])) - history["xrms"].append(xrms) - history["yrms"].append(yrms) + history["rms_x"].append(x_rms) + history["rms_y"].append(y_rms) + history["eps_x"].append(eps_x) + history["eps_y"].append(eps_y) - print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f}") + print( + f"turn={turn} xrms={x_rms:0.3f} yrms={y_rms:0.3f} epsx={eps_x:0.3f} epsy={eps_y:0.3f}" + ) particles_out = collect_bunch(bunch_out)["coords"] particles_out = particles_out[:, (0, 1, 2, 3)] @@ -200,7 +211,6 @@ def main(args: argparse.Namespace) -> None: print("TRACK SCUnifEllipse2D") results["SCUnifEllipse"] = track(lattice, bunch, nturns=args.nturns) - # Analysis # ------------------------------------------------------------------------------ @@ -211,13 +221,15 @@ def main(args: argparse.Namespace) -> None: # Print errors for key in results["SCUnifEllipse"]["history"]: - deltas = results["SCUnifEllipse"]["history"][key] - results["SC2p5D"]["history"][key] + deltas = ( + results["SCUnifEllipse"]["history"][key] - results["SC2p5D"]["history"][key] + ) print("key:", key) print("max_abs_delta:", np.max(np.abs(deltas))) print("avg_abs_delta:", np.mean(np.abs(deltas))) - # Plot rms bunch sizes - for key in ["xrms", "yrms"]: + # Plot rms size and emittance + for key in ["eps_x", "eps_y", "rms_x", "rms_y"]: fig, ax = plt.subplots(figsize=(5, 3)) for i, model in enumerate(["SC2p5D", "SCUnifEllipse"]): plot_kws = {} @@ -226,7 +238,7 @@ def main(args: argparse.Namespace) -> None: ax.plot(results[model]["history"][key], marker=".", label=model, **plot_kws) ax.set_ylim(0.0, ax.get_ylim()[1] * 2.0) ax.set_xlabel("Turn") - ax.set_ylabel("RMS [mm]") + ax.set_ylabel(key) ax.legend(loc="upper right") plt.savefig(os.path.join(output_dir, f"fig_{key}")) plt.close() @@ -235,29 +247,27 @@ def main(args: argparse.Namespace) -> None: particles = results["SC2p5D"]["particles"] xmax = 4.0 * np.std(particles, axis=0) limits = list(zip(-xmax, xmax)) + dims = ["x", "xp", "y", "yp"] labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]"] - # Plot x-x' - fig, axs = plt.subplots(figsize=(6, 3), ncols=2, sharex=True, sharey=True) - for ax, model in zip(axs, results): - particles = results[model]["particles"] - ax.hist2d(particles[:, 0], particles[:, 1], bins=64, range=[limits[0], limits[1]]) - ax.set_xlabel(labels[0]) - ax.set_ylabel(labels[1]) - ax.set_title(model) - plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) - plt.close() - - # Plot y-y' - fig, axs = plt.subplots(figsize=(6, 3), ncols=2, sharex=True, sharey=True) - for ax, model in zip(axs, results): - particles = results[model]["particles"] - ax.hist2d(particles[:, 2], particles[:, 3], bins=64, range=[limits[2], limits[3]]) - ax.set_xlabel(labels[2]) - ax.set_ylabel(labels[3]) - ax.set_title(model) - plt.savefig(os.path.join(output_dir, "fig_dist_y_yp")) - plt.close() + # Plot x-x', y-y', x-y + for axis in [(0, 1), (2, 3), (0, 2)]: + fig, axs = plt.subplots(figsize=(6, 3), ncols=2, sharex=True, sharey=True) + for ax, model in zip(axs, results): + particles = results[model]["particles"] + values, edges = np.histogramdd( + particles[:, axis], bins=args.plot_bins, range=[limits[k] for k in axis] + ) + if args.plot_mask: + values = np.ma.masked_equal(values, 0.0) + ax.pcolormesh(edges[0], edges[1], values.T) + ax.set_xlabel(labels[axis[0]]) + ax.set_ylabel(labels[axis[1]]) + ax.set_title(model) + plt.savefig( + os.path.join(output_dir, f"fig_dist_{dims[axis[0]]}_{dims[axis[1]]}") + ) + plt.close() # Plot corner for model in results: @@ -265,8 +275,9 @@ def main(args: argparse.Namespace) -> None: fig, axs = plot_corner( particles, limits=limits, - bins=64, + bins=args.plot_bins, labels=labels, + mask=args.plot_mask, ) plt.savefig(os.path.join(output_dir, f"fig_dist_corner_{model}")) plt.close() From e8c151d6653603e1cbb9fa409cf6b6d8a1a10207 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Sat, 5 Sep 2026 01:29:57 -0400 Subject: [PATCH 5/6] Add KV tracking test --- tests/py/orbit/test_sc_unif_ellipse.py | 147 +++++++++++++++++++++++++ 1 file changed, 147 insertions(+) diff --git a/tests/py/orbit/test_sc_unif_ellipse.py b/tests/py/orbit/test_sc_unif_ellipse.py index 99d5d4d7..5481d7a6 100644 --- a/tests/py/orbit/test_sc_unif_ellipse.py +++ b/tests/py/orbit/test_sc_unif_ellipse.py @@ -1,10 +1,22 @@ import math +import random +import time import pytest +from orbit.bunch_generators import KVDist2D +from orbit.bunch_generators import TwissContainer from orbit.core.bunch import Bunch +from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse2D from orbit.core.spacecharge import SpaceChargeCalcUnifEllipse +from orbit.space_charge.sc2p5d import SC2p5D_AccNode +from orbit.space_charge.sc2p5d import SCUnifEllipse2D_AccNode +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.space_charge.sc2p5d import setSCUnifEllipse2DAccNodes +from orbit.teapot import DriftTEAPOT +from orbit.teapot import QuadTEAPOT +from orbit.teapot import TEAPOT_Lattice def test_uniform_ellipsoid_calculator_names(): @@ -75,3 +87,138 @@ def test_uniform_ellipse_2d_kick_matches_envelope_formula(): assert ex == pytest.approx(cos_phi * ex_local) assert ey == pytest.approx(-sin_phi * ex_local) assert ez == 0.0 + + +def _make_fodo_lattice(name): + lattice = TEAPOT_Lattice(name) + for index in range(2): + nodes = ( + QuadTEAPOT(f"qf_{index}", length=0.20, kq=+1.0), + DriftTEAPOT(f"drift_1_{index}", length=0.40), + QuadTEAPOT(f"qd_{index}", length=0.40, kq=-1.0), + DriftTEAPOT(f"drift_2_{index}", length=0.40), + QuadTEAPOT(f"qf_end_{index}", length=0.20, kq=+1.0), + ) + for node in nodes: + node.setnParts(2) + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + lattice.addNode(node) + lattice.initialize() + return lattice + + +def _make_kv_bunch(n_particles=16_384): + bunch = Bunch() + bunch.mass(0.93827231) + bunch.charge(1.0) + bunch.getSyncParticle().kinEnergy(1.0) + + distribution = KVDist2D( + TwissContainer(alpha=0.0, beta=2.0, emittance=1.0e-6), + TwissContainer(alpha=0.0, beta=2.0, emittance=1.0e-6), + ) + random_state = random.getstate() + random.seed(1_234_567) + try: + for _ in range(n_particles // 2): + x, xp, y, yp = distribution.getCoordinates() + z = 0.10 * (random.random() - 0.5) + bunch.addParticle(x, xp, y, yp, z, 0.0) + bunch.addParticle(-x, -xp, -y, -yp, -z, 0.0) + finally: + random.setstate(random_state) + + bunch.macroSize(2.0e11 / bunch.getSizeGlobal()) + return bunch + + +def _rms(bunch, coordinate): + values = [getattr(bunch, coordinate)(index) for index in range(bunch.getSize())] + average = sum(values) / len(values) + return math.sqrt(sum((value - average) ** 2 for value in values) / len(values)) + + +def _relative_particle_difference(reference, candidate, coordinate): + reference_values = [ + getattr(reference, coordinate)(index) for index in range(reference.getSize()) + ] + candidate_values = [ + getattr(candidate, coordinate)(index) for index in range(candidate.getSize()) + ] + average = sum(reference_values) / len(reference_values) + scale = math.sqrt( + sum((value - average) ** 2 for value in reference_values) + / len(reference_values) + ) + difference = math.sqrt( + sum( + (reference_value - candidate_value) ** 2 + for reference_value, candidate_value in zip( + reference_values, candidate_values + ) + ) + / len(reference_values) + ) + return difference / scale + + +def test_uniform_ellipse_2d_matches_grid_solver_for_kv_beam(): + initial_bunch = _make_kv_bunch() + grid_bunch = Bunch() + ellipse_bunch = Bunch() + initial_bunch.copyBunchTo(grid_bunch) + initial_bunch.copyBunchTo(ellipse_bunch) + + grid_lattice = _make_fodo_lattice("grid") + ellipse_lattice = _make_fodo_lattice("uniform ellipse") + path_length_min = 0.10 + + grid_nodes = setSC2p5DAccNodes( + grid_lattice, + path_length_min, + SpaceChargeCalc2p5D(64, 64, 1), + ) + ellipse_nodes = setSCUnifEllipse2DAccNodes( + ellipse_lattice, + path_length_min, + SpaceChargeCalcUnifEllipse2D(), + ) + assert grid_nodes and all(isinstance(node, SC2p5D_AccNode) for node in grid_nodes) + assert ellipse_nodes and all( + isinstance(node, SCUnifEllipse2D_AccNode) for node in ellipse_nodes + ) + assert len(grid_nodes) == len(ellipse_nodes) + + start = time.perf_counter() + grid_lattice.trackBunch(grid_bunch) + grid_time = time.perf_counter() - start + start = time.perf_counter() + ellipse_lattice.trackBunch(ellipse_bunch) + ellipse_time = time.perf_counter() - start + + grid_rms = {coordinate: _rms(grid_bunch, coordinate) for coordinate in ("x", "y")} + ellipse_rms = { + coordinate: _rms(ellipse_bunch, coordinate) for coordinate in ("x", "y") + } + for coordinate in ("x", "y"): + assert ellipse_rms[coordinate] == pytest.approx( + grid_rms[coordinate], rel=0.02 + ) + + distribution_difference = { + coordinate: _relative_particle_difference( + grid_bunch, ellipse_bunch, coordinate + ) + for coordinate in ("x", "xp", "y", "yp") + } + assert all(value < 0.05 for value in distribution_difference.values()) + + print( + f"SC2p5D: {grid_time:.3f} s, " + f"SCUnifEllipse2D: {ellipse_time:.3f} s, " + f"rms: {grid_rms} vs. {ellipse_rms}, " + f"distribution difference: {distribution_difference}" + ) + +test_uniform_ellipse_2d_matches_grid_solver_for_kv_beam() \ No newline at end of file From cff8c5698c57971b98abd93226f07ed13a0bd771 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Sat, 5 Sep 2026 14:10:22 -0400 Subject: [PATCH 6/6] Copy particles --- examples/SpaceCharge/scUnifEllipse/track_fodo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/SpaceCharge/scUnifEllipse/track_fodo.py b/examples/SpaceCharge/scUnifEllipse/track_fodo.py index fac68557..bbf4ce28 100644 --- a/examples/SpaceCharge/scUnifEllipse/track_fodo.py +++ b/examples/SpaceCharge/scUnifEllipse/track_fodo.py @@ -175,7 +175,7 @@ def track(lattice: TEAPOT_Lattice, bunch: Bunch, nturns: int) -> dict: particles_out = particles_out[:, (0, 1, 2, 3)] particles_out *= 1000.0 return { - "particles": particles_out, + "particles": particles_out.copy(), "history": history, }