From 69fe7bea61ce5d23b6596f6217fe7a1ba47fc224 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 21 Jan 2026 14:59:31 -0500 Subject: [PATCH 01/12] Add empty files --- src/spacecharge/LSpaceChargeCalcDeriv.cc | 0 src/spacecharge/LSpaceChargeCalcDeriv.hh | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 src/spacecharge/LSpaceChargeCalcDeriv.cc create mode 100644 src/spacecharge/LSpaceChargeCalcDeriv.hh diff --git a/src/spacecharge/LSpaceChargeCalcDeriv.cc b/src/spacecharge/LSpaceChargeCalcDeriv.cc new file mode 100644 index 00000000..e69de29b diff --git a/src/spacecharge/LSpaceChargeCalcDeriv.hh b/src/spacecharge/LSpaceChargeCalcDeriv.hh new file mode 100644 index 00000000..e69de29b From f679bd7da73f763b43e619ca5bda52f36e9b54d1 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 2 Feb 2026 16:21:44 -0500 Subject: [PATCH 02/12] Add nFreq parameter to LSpaceChargeCalc --- src/spacecharge/LSpaceChargeCalc.cc | 3 ++- src/spacecharge/LSpaceChargeCalc.hh | 3 ++- src/spacecharge/wrap_lspacechargecalc.cc | 7 ++++--- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 12737969..57af126f 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -27,7 +27,7 @@ using namespace OrbitUtils; -LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int nBins_in): CppPyWrapper(NULL) +LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int nBins_in, int nFreq_in): CppPyWrapper(NULL) { b_a = b_a_in; length = length_in; @@ -35,6 +35,7 @@ LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosM useSpaceCharge = useSpaceCharge_in; nBins = nBins_in; zGrid = new Grid1D(nBins, length); + nFreq = nFreq_in; _fftmagnitude = new double[nBins / 2]; _fftphase = new double[nBins / 2]; diff --git a/src/spacecharge/LSpaceChargeCalc.hh b/src/spacecharge/LSpaceChargeCalc.hh index 24fc00a6..66cec712 100644 --- a/src/spacecharge/LSpaceChargeCalc.hh +++ b/src/spacecharge/LSpaceChargeCalc.hh @@ -29,7 +29,7 @@ class LSpaceChargeCalc: public OrbitUtils::CppPyWrapper public: /** Constructor */ - LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int zSize_in); + LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int zSize_in, int nFreq_in); /** Destructor */ virtual ~LSpaceChargeCalc(); @@ -52,6 +52,7 @@ public: int nBins; int nMacrosMin; int useSpaceCharge; + int nFreq; //protected: Grid1D* zGrid; diff --git a/src/spacecharge/wrap_lspacechargecalc.cc b/src/spacecharge/wrap_lspacechargecalc.cc index a3a98ea2..9439623b 100644 --- a/src/spacecharge/wrap_lspacechargecalc.cc +++ b/src/spacecharge/wrap_lspacechargecalc.cc @@ -46,12 +46,13 @@ extern "C" int nMacrosMin; int useSpaceCharge; int nBins; + int nFreq; - if(!PyArg_ParseTuple(args,"ddiii:arguments",&b_a,&length,&nMacrosMin,&useSpaceCharge,&nBins)){ - ORBIT_MPI_Finalize("PyLSpaceChargeCalc - LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins) - constructor needs parameters."); + if(!PyArg_ParseTuple(args,"ddiiii:arguments", &b_a, &length, &nMacrosMin, &useSpaceCharge, &nBins, &nFreq)){ + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins, nFreq) - constructor needs parameters."); } - self->cpp_obj = new LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins); + self->cpp_obj = new LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins, nFreq); ((LSpaceChargeCalc*) self->cpp_obj)->setPyWrapper((PyObject*) self); From 69d66de070324d5d9907e1b212f2a81bf2e11708 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 2 Feb 2026 16:33:00 -0500 Subject: [PATCH 03/12] Set number of frequency components --- src/spacecharge/LSpaceChargeCalc.cc | 59 ++++++++++++++--------------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 57af126f..79dbcdf5 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -35,7 +35,14 @@ LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosM useSpaceCharge = useSpaceCharge_in; nBins = nBins_in; zGrid = new Grid1D(nBins, length); + nFreq = nFreq_in; + if (nFreq > (nBins / 2)) { + nFreq = nBins / 2; + } + if (nFreq < 0) { + nFreq = nBins / 2; + } _fftmagnitude = new double[nBins / 2]; _fftphase = new double[nBins / 2]; @@ -92,11 +99,14 @@ void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) void LSpaceChargeCalc::trackBunch(Bunch* bunch) { + + std::cout << "nbins=" << nBins << " nfreq=" << nFreq << std::endl; + + int nPartsGlobal = bunch->getSizeGlobal(); if(nPartsGlobal < nMacrosMin) return; -// Bin the particles - + // Bin the particles double zmin, zmax; double realPart, imagPart; @@ -109,8 +119,7 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) zGrid->binBunchSmoothedByParticle(bunch); zGrid->synchronizeMPI(bunch->getMPI_Comm_Local()); -// FFT the beam density - + // FFT the beam density for(int i = 0; i < nBins; i++) { _in[i][0] = zGrid->getValueOnGrid(i); @@ -119,8 +128,7 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) fftw_execute(_plan); -// Find the magnitude and phase - + // Find the magnitude and phase _fftmagnitude[0] = _out[0][0]/(double)nBins; _fftphase[0] = 0.; @@ -132,16 +140,13 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) _fftphase[n] = atan2(imagPart,realPart); } -// Space charge contribution to impedance - + // Space charge contribution to impedance SyncPart* sp = bunch->getSyncPart(); -// Set zero if useSpaceCharge = 0 - + // Set zero if useSpaceCharge = 0 double zSpaceCharge_n = 0.; -// Otherwise, set positive since space charge is capacitive (Chao convention) - + // Otherwise, set positive since space charge is capacitive (Chao convention) if(useSpaceCharge != 0) { double mu_0 = 4.0 * OrbitConst::PI * 1.e-07; // permeability of free space @@ -159,17 +164,13 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) _chi[n] = atan2(imagPart, realPart); } -// Convert charge to current for a single macroparticle per unit bin length - - double charge2current = bunch->getCharge() * bunch->getMacroSize() * - OrbitConst::elementary_charge_MKS * sp->getBeta() * - OrbitConst::c / (length / nBins); - -// Calculate and add the kick to macroparticles + // Convert charge to current for a single macroparticle per unit bin length + double charge2current = bunch->getCharge() * bunch->getMacroSize() * OrbitConst::elementary_charge_MKS * sp->getBeta() * OrbitConst::c / (length / nBins); + // Calculate and add the kick to macroparticles double kick, angle, position; -// Convert particle longitudinal coordinate to phi + // Convert particle longitudinal coordinate to phi double philocal; double z; @@ -179,14 +180,12 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) z = bunch->z(j); philocal = (z / length) * 2 * OrbitConst::PI; -// Handle cases where the longitudinal coordinate is -// outside of the user-specified length - + // Handle cases where the longitudinal coordinate is + // outside of the user-specified length if(philocal < -OrbitConst::PI) philocal += 2 * OrbitConst::PI; if(philocal > OrbitConst::PI) philocal -= 2 * OrbitConst::PI; - double dE = _kick(philocal) * (-1e-9) * - bunch->getCharge() * charge2current; + double dE = _kick(philocal) * (-1e-9) * bunch->getCharge() * charge2current; coords[j][5] += dE; } } @@ -211,15 +210,15 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) double LSpaceChargeCalc::_kick(double angle) { -// n=0 term has no impact (constant in phi) -// f(phi) = _FFTMagnitude(1) + sum (n = 2 -> N / 2) of -// [2 * _FFTMagnitude(i) * cos(phi * (n - 1) + -// _FFTPhase(n) + _chi(n))] + // n=0 term has no impact (constant in phi) + // f(phi) = _FFTMagnitude(1) + sum (n = 2 -> N / 2) of + // [2 * _FFTMagnitude(i) * cos(phi * (n - 1) + + // _FFTPhase(n) + _chi(n))] double kick = 0.; double cosArg; - for(int n = 1; n < nBins / 2; n++) + for(int n = 1; n < nFreq; n++) { cosArg = n * (angle + OrbitConst::PI) + _fftphase[n] + _chi[n]; kick += 2 * _fftmagnitude[n] * _z[n] * cos(cosArg); From 1e5f1cfc1503399b58a2e5686a88772b4dbaf82e Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 2 Feb 2026 21:39:45 -0500 Subject: [PATCH 04/12] Remove print statement --- src/spacecharge/LSpaceChargeCalc.cc | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 79dbcdf5..77fd91a4 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -99,10 +99,6 @@ void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) void LSpaceChargeCalc::trackBunch(Bunch* bunch) { - - std::cout << "nbins=" << nBins << " nfreq=" << nFreq << std::endl; - - int nPartsGlobal = bunch->getSizeGlobal(); if(nPartsGlobal < nMacrosMin) return; From d3dcbd7f6295a14af6f8a4b6041ed81efbecd78c Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 2 Feb 2026 21:40:16 -0500 Subject: [PATCH 05/12] Add nfreq option to LSpaceChargeNode --- py/orbit/space_charge/sc1d/sc1DNode.py | 42 ++++++++------------------ 1 file changed, 13 insertions(+), 29 deletions(-) diff --git a/py/orbit/space_charge/sc1d/sc1DNode.py b/py/orbit/space_charge/sc1d/sc1DNode.py index 628eaf53..b8b876ae 100644 --- a/py/orbit/space_charge/sc1d/sc1DNode.py +++ b/py/orbit/space_charge/sc1d/sc1DNode.py @@ -6,33 +6,25 @@ import os import math -# import the function that finalizes the execution -from orbit.utils import orbitFinalize - -# import physical constants +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import LSpaceChargeCalc +from orbit.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.lattice import AccActionsContainer +from orbit.lattice import AccNodeBunchTracker from orbit.utils import consts - -# import general accelerator elements and lattice -from orbit.lattice import AccLattice, AccNode, AccActionsContainer, AccNodeBunchTracker - -# import teapot drift class +from orbit.utils import orbitFinalize from orbit.teapot import DriftTEAPOT -# import longitudinal space charge package -from orbit.core.spacecharge import LSpaceChargeCalc - -# ----------------------------------------------------------------------------- -# Node for impedance as function of node number -# ----------------------------------------------------------------------------- - class SC1D_AccNode(DriftTEAPOT): - def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, name="long sc node"): + """Longitudinal space charge node.""" + def __init__(self, b_a: float, phase_length: float, nmacros_min: float, use_sc: float, nbins: float, nfreq: int = -1, name="long sc node"): """ Constructor. Creates the SC1D-teapot element. """ DriftTEAPOT.__init__(self, name) - self.lspacecharge = LSpaceChargeCalc(b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins) + self.lspacecharge = LSpaceChargeCalc(b_a, phase_length, nmacros_min, use_sc, nbins, nfreq) self.setType("long sc node") self.setLength(0.0) @@ -51,18 +43,14 @@ def track(self, paramsDict): """ length = self.getLength(self.getActivePartIndex()) bunch = paramsDict["bunch"] - self.lspacecharge.trackBunch(bunch) # track method goes here + self.lspacecharge.trackBunch(bunch) def assignImpedance(self, py_cmplx_arr): self.lspacecharge.assignImpedance(py_cmplx_arr) -# ----------------------------------------------------------------------------- -# Node for impedance as function of frequency -# ----------------------------------------------------------------------------- - - class FreqDep_SC1D_AccNode(DriftTEAPOT): + """Longitudinal space charge node (frequency-dependent).""" def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, impeDict, name="freq. dep. long sc node"): """ Constructor. Creates the FreqDep_SC1D-teapot element. @@ -121,12 +109,8 @@ def track(self, paramsDict): self.lspacecharge.trackBunch(bunch) -# ----------------------------------------------------------------------------- -# Node for impedance as function of beta and frequency -# ----------------------------------------------------------------------------- - - class BetFreqDep_SC1D_AccNode(DriftTEAPOT): + """Longitudinal space charge node (frequency- and velocity-dependent).""" def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, impeDict, name="freq. dep. long sc node"): """ Constructor. Creates the BetFreqDep_SC1D-teapot element. From d461d355967288ac2581d8dcb73f58e6b23d7267 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 2 Feb 2026 21:40:45 -0500 Subject: [PATCH 06/12] Remove empty files --- src/spacecharge/LSpaceChargeCalcDeriv.cc | 0 src/spacecharge/LSpaceChargeCalcDeriv.hh | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 src/spacecharge/LSpaceChargeCalcDeriv.cc delete mode 100644 src/spacecharge/LSpaceChargeCalcDeriv.hh diff --git a/src/spacecharge/LSpaceChargeCalcDeriv.cc b/src/spacecharge/LSpaceChargeCalcDeriv.cc deleted file mode 100644 index e69de29b..00000000 diff --git a/src/spacecharge/LSpaceChargeCalcDeriv.hh b/src/spacecharge/LSpaceChargeCalcDeriv.hh deleted file mode 100644 index e69de29b..00000000 From c8e92e3acfc5ef6d6986748d295fed8b14df672a Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Tue, 3 Feb 2026 13:15:13 -0500 Subject: [PATCH 07/12] Set number of FFT modes used in longitudinal space charge kick --- src/spacecharge/LSpaceChargeCalc.cc | 31 ++++++++++++++------ src/spacecharge/LSpaceChargeCalc.hh | 17 +++++++---- src/spacecharge/wrap_lspacechargecalc.cc | 37 +++++++++++++++++++++--- 3 files changed, 66 insertions(+), 19 deletions(-) diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 77fd91a4..56dc99a9 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -27,7 +27,7 @@ using namespace OrbitUtils; -LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int nBins_in, int nFreq_in): CppPyWrapper(NULL) +LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int nBins_in): CppPyWrapper(NULL) { b_a = b_a_in; length = length_in; @@ -36,13 +36,8 @@ LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosM nBins = nBins_in; zGrid = new Grid1D(nBins, length); - nFreq = nFreq_in; - if (nFreq > (nBins / 2)) { - nFreq = nBins / 2; - } - if (nFreq < 0) { - nFreq = nBins / 2; - } + nModes = nBins / 2; + useGrad = 0; _fftmagnitude = new double[nBins / 2]; _fftphase = new double[nBins / 2]; @@ -91,6 +86,24 @@ LSpaceChargeCalc::~LSpaceChargeCalc() } +void LSpaceChargeCalc::setNumModes(int n) { + nModes = n; + int nModesMax = nBins / 2; + + if (nModes > nModesMax) { + nModes = nModesMax; + } + if (nModes < 0) { + nModes = nModesMax; + } +} + + +void LSpaceChargeCalc::setUseGrad(int setting) { + useGrad = setting; +} + + void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) { _zImped_n[n + 1] = std::complex(real, imag); @@ -214,7 +227,7 @@ double LSpaceChargeCalc::_kick(double angle) double kick = 0.; double cosArg; - for(int n = 1; n < nFreq; n++) + for(int n = 1; n < nModes; n++) { cosArg = n * (angle + OrbitConst::PI) + _fftphase[n] + _chi[n]; kick += 2 * _fftmagnitude[n] * _z[n] * cos(cosArg); diff --git a/src/spacecharge/LSpaceChargeCalc.hh b/src/spacecharge/LSpaceChargeCalc.hh index 66cec712..4e81c9e2 100644 --- a/src/spacecharge/LSpaceChargeCalc.hh +++ b/src/spacecharge/LSpaceChargeCalc.hh @@ -29,17 +29,21 @@ class LSpaceChargeCalc: public OrbitUtils::CppPyWrapper public: /** Constructor */ - LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int zSize_in, int nFreq_in); + LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int zSize_in); /** Destructor */ virtual ~LSpaceChargeCalc(); - /** Calculates space charge and applies the transverse and - longitudinal SC kicks to the macro-particles in the bunch. */ + /** Sets number of FFT modes used to calculate energy kick from impedance. **/ + void setNumModes(int n); + + /** Sets option to use gradient rather than impedance. **/ + void setUseGrad(int n); + + /** Calculates space charge and applies the transverse and longitudinal SC kicks to the macro-particles in the bunch. */ void trackBunch(Bunch* bunch); - /** Assigns the real and imaginary parts of the - machine impedance for index n**/ + /** Assigns the real and imaginary parts of the machine impedance for index n**/ void assignImpedanceValue(int n, double real, double imag); /** Routine for calculating the kick to the particle **/ @@ -52,7 +56,8 @@ public: int nBins; int nMacrosMin; int useSpaceCharge; - int nFreq; + int nModes; + int useGrad; //protected: Grid1D* zGrid; diff --git a/src/spacecharge/wrap_lspacechargecalc.cc b/src/spacecharge/wrap_lspacechargecalc.cc index 9439623b..6e04f98d 100644 --- a/src/spacecharge/wrap_lspacechargecalc.cc +++ b/src/spacecharge/wrap_lspacechargecalc.cc @@ -46,13 +46,12 @@ extern "C" int nMacrosMin; int useSpaceCharge; int nBins; - int nFreq; - if(!PyArg_ParseTuple(args,"ddiiii:arguments", &b_a, &length, &nMacrosMin, &useSpaceCharge, &nBins, &nFreq)){ - ORBIT_MPI_Finalize("PyLSpaceChargeCalc - LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins, nFreq) - constructor needs parameters."); + if(!PyArg_ParseTuple(args,"ddiiii:arguments", &b_a, &length, &nMacrosMin, &useSpaceCharge, &nBins)){ + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins) - constructor needs parameters."); } - self->cpp_obj = new LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins, nFreq); + self->cpp_obj = new LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins); ((LSpaceChargeCalc*) self->cpp_obj)->setPyWrapper((PyObject*) self); @@ -93,7 +92,37 @@ extern "C" } + static PyObject* LSpaceChargeCalc_setNumModes(PyObject *self, PyObject *args){ + pyORBIT_Object* pyLSpaceChargeCalc = (pyORBIT_Object*) self; + LSpaceChargeCalc* cpp_LSpaceChargeCalc = (LSpaceChargeCalc*) pyLSpaceChargeCalc->cpp_obj; + + int n = -1; + + if(!PyArg_ParseTuple(args,"i:arguments", &n)){ + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setNumModes(n) - constructor needs parameters."); + } + cpp_LSpaceChargeCalc->setNumModes(n); + + Py_INCREF(Py_None); + return Py_None; + } + + static PyObject* LSpaceChargeCalc_setUseGrad(PyObject *self, PyObject *args){ + + pyORBIT_Object* pyLSpaceChargeCalc = (pyORBIT_Object*) self; + LSpaceChargeCalc* cpp_LSpaceChargeCalc = (LSpaceChargeCalc*) pyLSpaceChargeCalc->cpp_obj; + + int setting = 0; + + if(!PyArg_ParseTuple(args,"i:arguments", &setting)){ + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setUseGrad(setting) - constructor needs parameters."); + } + cpp_LSpaceChargeCalc->setUseGrad(setting); + + Py_INCREF(Py_None); + return Py_None; + } //assignImpedanceValue(int, real, real). Wraps the LongSpaceChargeCalc routine assigning an impedance mode static PyObject* LSpaceChargeCalc_assignImpedanceValue(PyObject *self, PyObject *args){ From e372ff8411ad03f841a18e87f818d3e31c13b533 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Tue, 3 Feb 2026 14:01:59 -0500 Subject: [PATCH 08/12] Save progress; getting error with method setNumModes --- py/orbit/space_charge/sc1d/sc1DNode.py | 157 +++++++--- src/spacecharge/LSpaceChargeCalc.cc | 348 +++++++++++------------ src/spacecharge/LSpaceChargeCalc.hh | 2 +- src/spacecharge/wrap_lspacechargecalc.cc | 52 ++-- 4 files changed, 311 insertions(+), 248 deletions(-) diff --git a/py/orbit/space_charge/sc1d/sc1DNode.py b/py/orbit/space_charge/sc1d/sc1DNode.py index b8b876ae..0bf968e1 100644 --- a/py/orbit/space_charge/sc1d/sc1DNode.py +++ b/py/orbit/space_charge/sc1d/sc1DNode.py @@ -19,16 +19,37 @@ class SC1D_AccNode(DriftTEAPOT): """Longitudinal space charge node.""" - def __init__(self, b_a: float, phase_length: float, nmacros_min: float, use_sc: float, nbins: float, nfreq: int = -1, name="long sc node"): + + def __init__( + self, + b_a: float, + phase_length: float, + nmacros_min: float, + use_sc: float, + nbins: float, + nmodes: int = None, + use_grad: bool = False, + name="long sc node", + ) -> None: """ Constructor. Creates the SC1D-teapot element. """ DriftTEAPOT.__init__(self, name) - self.lspacecharge = LSpaceChargeCalc(b_a, phase_length, nmacros_min, use_sc, nbins, nfreq) + self.lspacecharge = LSpaceChargeCalc(b_a, phase_length, nmacros_min, use_sc, nbins) + self.setNumModes(nmodes) + # self.setUseGrad(use_grad) self.setType("long sc node") self.setLength(0.0) - def trackBunch(self, bunch): + def setUseGrad(self, use_grad: bool) -> None: + """Sets whether to use gradient-based solver instead of impedance solver.""" + self.lspacecharge.setUseGrad(int(use_grad)) + + def setNumModes(self, n: int) -> None: + """Sets number of FFT modes used to calculate energy kick.""" + self.lspacecharge.setNumModes(n) + + def trackBunch(self, bunch: Bunch) -> None: """ The SC1D-teapot class implementation of the AccNodeBunchTracker class trackBunch(probe) method. @@ -36,72 +57,85 @@ def trackBunch(self, bunch): length = self.getLength(self.getActivePartIndex()) self.lspacecharge.trackBunch(bunch) # track method goes here - def track(self, paramsDict): + def track(self, params_dict: dict) -> None: """ The SC1D-teapot class implementation of the AccNodeBunchTracker class track(probe) method. """ length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] + bunch = params_dict["bunch"] self.lspacecharge.trackBunch(bunch) - def assignImpedance(self, py_cmplx_arr): + def assignImpedance(self, py_cmplx_arr: list[float]) -> None: self.lspacecharge.assignImpedance(py_cmplx_arr) class FreqDep_SC1D_AccNode(DriftTEAPOT): """Longitudinal space charge node (frequency-dependent).""" - def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, impeDict, name="freq. dep. long sc node"): + + def __init__( + self, + b_a: float, + phase_length: float, + nmacros_min: int, + use_sc: int, + nbins: int, + bunch: Bunch, + imp_dict: dict, + name: str = "freq. dep. long sc node", + ) -> None: """ Constructor. Creates the FreqDep_SC1D-teapot element. """ DriftTEAPOT.__init__(self, name) - self.lspacecharge = LSpaceChargeCalc(b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins) + self.lspacecharge = LSpaceChargeCalc( + b_a, phase_length, nmacros_min, use_sc, nbins + ) self.setType("freq. dep. long sc node") self.setLength(0.0) - self.phaseLength = phaseLength - self.nBins = nBins - self.localDict = impeDict + self.phase_length = phase_length + self.nbins = nbins + self.localDict = imp_dict self.freq_tuple = self.localDict["freqs"] self.freq_range = len(self.freq_tuple) - 1 self.z_tuple = self.localDict["z_imp"] self.c = consts.speed_of_light BetaRel = bunch.getSyncParticle().beta() - Freq0 = (BetaRel * self.c) / self.phaseLength + Freq0 = (BetaRel * self.c) / self.phase_length Z = [] - for n in range(self.nBins // 2 - 1): + for n in range(self.nbins // 2 - 1): freq_mode = Freq0 * (n + 1) z_mode = interp(freq_mode, self.freq_range, self.freq_tuple, self.z_tuple) Z.append(z_mode) self.lspacecharge.assignImpedance(Z) - def trackBunch(self, bunch): + def trackBunch(self, bunch: Bunch) -> None: """ The FreqDep_SC1D-teapot class implementation of the AccNodeBunchTracker class track(probe) method. """ length = self.getLength(self.getActivePartIndex()) BetaRel = bunch.getSyncParticle().beta() - Freq0 = (BetaRel * self.c) / self.phaseLength + Freq0 = (BetaRel * self.c) / self.phase_length Z = [] - for n in range(self.nBins // 2 - 1): + for n in range(self.nbins // 2 - 1): freq_mode = Freq0 * (n + 1) z_mode = interp(freq_mode, self.freq_range, self.freq_tuple, self.z_tuple) Z.append(z_mode) self.lspacecharge.assignImpedance(Z) self.lspacecharge.trackBunch(bunch) - def track(self, paramsDict): + def track(self, params_dict: dict) -> None: """ The FreqDep_SC1D-teapot class implementation of the AccNodeBunchTracker class track(probe) method. """ length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] + bunch = params_dict["bunch"] BetaRel = bunch.getSyncParticle().beta() - Freq0 = (BetaRel * self.c) / self.phaseLength + Freq0 = (BetaRel * self.c) / self.phase_length Z = [] - for n in range(self.nBins // 2 - 1): + for n in range(self.nbins // 2 - 1): freq_mode = Freq0 * (n + 1) z_mode = interp(freq_mode, self.freq_range, self.freq_tuple, self.z_tuple) Z.append(z_mode) @@ -111,17 +145,30 @@ def track(self, paramsDict): class BetFreqDep_SC1D_AccNode(DriftTEAPOT): """Longitudinal space charge node (frequency- and velocity-dependent).""" - def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, impeDict, name="freq. dep. long sc node"): + + def __init__( + self, + b_a: float, + phase_length: float, + nmacros_min: float, + use_sc: int, + nbins: int, + bunch: Bunch, + imp_dict: dict, + name: str = "freq. dep. long sc node", + ) -> None: """ Constructor. Creates the BetFreqDep_SC1D-teapot element. """ DriftTEAPOT.__init__(self, name) - self.lspacecharge = LSpaceChargeCalc(b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins) + self.lspacecharge = LSpaceChargeCalc( + b_a, phase_length, nmacros_min, use_sc, nbins + ) self.setType("beta-freq. dep. long sc node") self.setLength(0.0) - self.phaseLength = phaseLength - self.nBins = nBins - self.localDict = impeDict + self.phase_length = phase_length + self.nbins = nbins + self.localDict = imp_dict self.bet_tuple = self.localDict["betas"] self.bet_range = len(self.bet_tuple) - 1 self.freq_tuple = self.localDict["freqs"] @@ -129,49 +176,73 @@ def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, i self.z_bf = self.localDict["z_imp"] self.c = consts.speed_of_light BetaRel = bunch.getSyncParticle().beta() - Freq0 = (BetaRel * self.c) / self.phaseLength + Freq0 = (BetaRel * self.c) / self.phase_length Z = [] - for n in range(self.nBins / 2 - 1): + for n in range(self.nbins / 2 - 1): freq_mode = Freq0 * (n + 1) - z_mode = bilinterp(BetaRel, freq_mode, self.bet_range, self.freq_range, self.bet_tuple, self.freq_tuple, self.z_bf) + z_mode = bilinterp( + BetaRel, + freq_mode, + self.bet_range, + self.freq_range, + self.bet_tuple, + self.freq_tuple, + self.z_bf, + ) Z.append(z_mode) self.lspacecharge.assignImpedance(Z) - def trackBunch(self, bunch): + def trackBunch(self, bunch: Bunch) -> None: """ The BetFreqDep_SC1D-teapot class implementation of the AccNodeBunchTracker class track(probe) method. """ length = self.getLength(self.getActivePartIndex()) BetaRel = bunch.getSyncParticle().beta() - Freq0 = (BetaRel * self.c) / self.phaseLength + Freq0 = (BetaRel * self.c) / self.phase_length Z = [] - for n in range(self.nBins / 2 - 1): + for n in range(self.nbins / 2 - 1): freq_mode = Freq0 * (n + 1) - z_mode = bilinterp(BetaRel, freq_mode, self.bet_range, self.freq_range, self.bet_tuple, self.freq_tuple, self.z_bf) + z_mode = bilinterp( + BetaRel, + freq_mode, + self.bet_range, + self.freq_range, + self.bet_tuple, + self.freq_tuple, + self.z_bf, + ) Z.append(z_mode) self.lspacecharge.assignImpedance(Z) self.lspacecharge.trackBunch(bunch) - def track(self, paramsDict): + def track(self, params_dict: dict) -> None: """ The BetFreqDep_SC1D-teapot class implementation of the AccNodeBunchTracker class track(probe) method. """ length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] + bunch = params_dict["bunch"] BetaRel = bunch.getSyncParticle().beta() - Freq0 = (BetaRel * self.c) / self.phaseLength + Freq0 = (BetaRel * self.c) / self.phase_length Z = [] - for n in range(self.nBins / 2 - 1): + for n in range(self.nbins / 2 - 1): freq_mode = Freq0 * (n + 1) - z_mode = bilinterp(BetaRel, freq_mode, self.bet_range, self.freq_range, self.bet_tuple, self.freq_tuple, self.z_bf) + z_mode = bilinterp( + BetaRel, + freq_mode, + self.bet_range, + self.freq_range, + self.bet_tuple, + self.freq_tuple, + self.z_bf, + ) Z.append(z_mode) self.lspacecharge.assignImpedance(Z) self.lspacecharge.trackBunch(bunch) -def interp(x, n_tuple, x_tuple, y_tuple): +def interp(x: float, n_tuple: int, x_tuple: list[float], y_tuple: list[float]) -> float: """ Linear interpolation: Given n-tuple + 1 points, x_tuple and y_tuple, routine finds y = y_tuple @@ -194,7 +265,15 @@ def interp(x, n_tuple, x_tuple, y_tuple): return y -def bilinterp(x, y, nx_tuple, ny_tuple, x_tuple, y_tuple, fxy): +def bilinterp( + x: float, + y: float, + nx_tuple: int, + ny_tuple: int, + x_tuple: list[float], + y_tuple: list[float], + fxy: list[list[float]], +) -> float: """ Bilinear interpolation: Given nx-tuple + 1 x-points, ny-tuple + 1 y-points, x_tuple and y_tuple, diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 56dc99a9..a24f49d5 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -11,195 +11,183 @@ // ///////////////////////////////////////////////////////////////////////////// -#include "Grid1D.hh" -#include "BufferStore.hh" #include "LSpaceChargeCalc.hh" +#include "BufferStore.hh" +#include "Grid1D.hh" #include "OrbitConst.hh" -#include -#include -#include #include +#include #include +#include -//FFTW library header +// FFTW library header #include "fftw3.h" using namespace OrbitUtils; - -LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int nBins_in): CppPyWrapper(NULL) -{ - b_a = b_a_in; - length = length_in; - nMacrosMin = nMacrosMin_in; - useSpaceCharge = useSpaceCharge_in; - nBins = nBins_in; - zGrid = new Grid1D(nBins, length); - - nModes = nBins / 2; - useGrad = 0; - - _fftmagnitude = new double[nBins / 2]; - _fftphase = new double[nBins / 2]; - _z = new double[nBins / 2]; - _chi = new double[nBins / 2]; - //--------------------------- - // _zImped_n size was increased by 1 to avoid crash due to - // assignment in LSpaceChargeCalc::assignImpedanceValue(...) - // to _zImped_n[n + 1] element. This should be resolved later. - //---------------------------- - _zImped_n = new std::complex[nBins / 2 + 1]; - - for(int n = 0; n < nBins / 2; n++) - { - _fftmagnitude[n] = 0.0; - _fftphase[n] = 0.0; - _z[n] = 0.0; - _chi[n] = 0.0; - _zImped_n[n] = std::complex(0.0, 0.0); - } - - _in = (fftw_complex *) fftw_malloc(nBins * sizeof(fftw_complex)); - _out = (fftw_complex *) fftw_malloc(nBins * sizeof(fftw_complex)); - _plan = fftw_plan_dft_1d(nBins, _in, _out, FFTW_FORWARD, FFTW_ESTIMATE); +LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosMin_in, int useSpaceCharge_in, int nBins_in) : CppPyWrapper(NULL) { + b_a = b_a_in; + length = length_in; + nMacrosMin = nMacrosMin_in; + useSpaceCharge = useSpaceCharge_in; + nBins = nBins_in; + zGrid = new Grid1D(nBins, length); + + nModes = nBins / 2; + useGrad = 0; + + _fftmagnitude = new double[nBins / 2]; + _fftphase = new double[nBins / 2]; + _z = new double[nBins / 2]; + _chi = new double[nBins / 2]; + //--------------------------- + // _zImped_n size was increased by 1 to avoid crash due to + // assignment in LSpaceChargeCalc::assignImpedanceValue(...) + // to _zImped_n[n + 1] element. This should be resolved later. + //---------------------------- + _zImped_n = new std::complex[nBins / 2 + 1]; + + for (int n = 0; n < nBins / 2; n++) { + _fftmagnitude[n] = 0.0; + _fftphase[n] = 0.0; + _z[n] = 0.0; + _chi[n] = 0.0; + _zImped_n[n] = std::complex(0.0, 0.0); + } + + _in = (fftw_complex *)fftw_malloc(nBins * sizeof(fftw_complex)); + _out = (fftw_complex *)fftw_malloc(nBins * sizeof(fftw_complex)); + _plan = fftw_plan_dft_1d(nBins, _in, _out, FFTW_FORWARD, FFTW_ESTIMATE); } - -LSpaceChargeCalc::~LSpaceChargeCalc() -{ - if(zGrid->getPyWrapper() != NULL) - { - Py_DECREF(zGrid->getPyWrapper()); - } - else - { - delete zGrid; - } - delete[] _fftmagnitude; - delete[] _fftphase; - delete[] _z; - delete[] _chi; - delete[] _zImped_n; - fftw_free(_in); - fftw_free(_out); - fftw_destroy_plan(_plan); +LSpaceChargeCalc::~LSpaceChargeCalc() { + if (zGrid->getPyWrapper() != NULL) { + Py_DECREF(zGrid->getPyWrapper()); + } else { + delete zGrid; + } + delete[] _fftmagnitude; + delete[] _fftphase; + delete[] _z; + delete[] _chi; + delete[] _zImped_n; + fftw_free(_in); + fftw_free(_out); + fftw_destroy_plan(_plan); } - void LSpaceChargeCalc::setNumModes(int n) { - nModes = n; - int nModesMax = nBins / 2; - - if (nModes > nModesMax) { - nModes = nModesMax; - } - if (nModes < 0) { - nModes = nModesMax; - } + nModes = n; + int nModesMax = nBins / 2; + + if (nModes > nModesMax) { + nModes = nModesMax; + } + if (nModes < 0) { + nModes = nModesMax; + } } - void LSpaceChargeCalc::setUseGrad(int setting) { useGrad = setting; } - -void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) -{ - _zImped_n[n + 1] = std::complex(real, imag); +void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) { + _zImped_n[n + 1] = std::complex(real, imag); } - -void LSpaceChargeCalc::trackBunch(Bunch* bunch) -{ - int nPartsGlobal = bunch->getSizeGlobal(); - if(nPartsGlobal < nMacrosMin) return; - - // Bin the particles - double zmin, zmax; - double realPart, imagPart; - - bunchExtremaCalc->getExtremaZ(bunch, zmin, zmax); - double zextra = (length - (zmax - zmin)) / 2.0; - zmax += zextra; - zmin = zmax - length; - zGrid->setGridZ(zmin, zmax); - zGrid->setZero(); - zGrid->binBunchSmoothedByParticle(bunch); - zGrid->synchronizeMPI(bunch->getMPI_Comm_Local()); - - // FFT the beam density - for(int i = 0; i < nBins; i++) - { - _in[i][0] = zGrid->getValueOnGrid(i); - _in[i][1] = 0.; - } - - fftw_execute(_plan); - - // Find the magnitude and phase - _fftmagnitude[0] = _out[0][0]/(double)nBins; - _fftphase[0] = 0.; - - for(int n = 1; n < nBins / 2; n++) - { - realPart = _out[n][0]/((double)nBins); - imagPart = _out[n][1]/((double)nBins); - _fftmagnitude[n] = sqrt(realPart * realPart + imagPart * imagPart); - _fftphase[n] = atan2(imagPart,realPart); - } - - // Space charge contribution to impedance - SyncPart* sp = bunch->getSyncPart(); - - // Set zero if useSpaceCharge = 0 - double zSpaceCharge_n = 0.; - - // Otherwise, set positive since space charge is capacitive (Chao convention) - if(useSpaceCharge != 0) - { - double mu_0 = 4.0 * OrbitConst::PI * 1.e-07; // permeability of free space - double _z_0 = OrbitConst::c * mu_0; - - zSpaceCharge_n = _z_0 * (1.0 + 2.0 * log(b_a)) / - (2 * sp->getBeta() * sp->getGamma() * sp->getGamma()); - } - - for(int n = 1; n < nBins / 2; n++) - { - realPart = std::real(_zImped_n[n]); - imagPart = std::imag(_zImped_n[n]) + zSpaceCharge_n; - _z[n] = n * sqrt(realPart * realPart + imagPart * imagPart); - _chi[n] = atan2(imagPart, realPart); - } - - // Convert charge to current for a single macroparticle per unit bin length - double charge2current = bunch->getCharge() * bunch->getMacroSize() * OrbitConst::elementary_charge_MKS * sp->getBeta() * OrbitConst::c / (length / nBins); - - // Calculate and add the kick to macroparticles - double kick, angle, position; - - // Convert particle longitudinal coordinate to phi - - double philocal; - double z; - double** coords = bunch->coordArr(); - for (int j = 0; j < bunch->getSize(); j++) - { - z = bunch->z(j); - philocal = (z / length) * 2 * OrbitConst::PI; - - // Handle cases where the longitudinal coordinate is - // outside of the user-specified length - if(philocal < -OrbitConst::PI) philocal += 2 * OrbitConst::PI; - if(philocal > OrbitConst::PI) philocal -= 2 * OrbitConst::PI; - - double dE = _kick(philocal) * (-1e-9) * bunch->getCharge() * charge2current; - coords[j][5] += dE; - } +void LSpaceChargeCalc::trackBunch(Bunch *bunch) { + int nPartsGlobal = bunch->getSizeGlobal(); + if (nPartsGlobal < nMacrosMin) + return; + + // Bin the particles + double zmin, zmax; + double realPart, imagPart; + + bunchExtremaCalc->getExtremaZ(bunch, zmin, zmax); + double zextra = (length - (zmax - zmin)) / 2.0; + zmax += zextra; + zmin = zmax - length; + zGrid->setGridZ(zmin, zmax); + zGrid->setZero(); + zGrid->binBunchSmoothedByParticle(bunch); + zGrid->synchronizeMPI(bunch->getMPI_Comm_Local()); + + if (useGrad == 0) { + + // FFT the beam density + for (int i = 0; i < nBins; i++) { + _in[i][0] = zGrid->getValueOnGrid(i); + _in[i][1] = 0.; + } + + fftw_execute(_plan); + + // Find the magnitude and phase + _fftmagnitude[0] = _out[0][0] / (double)nBins; + _fftphase[0] = 0.; + + for (int n = 1; n < nBins / 2; n++) { + realPart = _out[n][0] / ((double)nBins); + imagPart = _out[n][1] / ((double)nBins); + _fftmagnitude[n] = sqrt(realPart * realPart + imagPart * imagPart); + _fftphase[n] = atan2(imagPart, realPart); + } + + // Space charge contribution to impedance + SyncPart *sp = bunch->getSyncPart(); + + // Set zero if useSpaceCharge = 0 + double zSpaceCharge_n = 0.; + + // Otherwise, set positive since space charge is capacitive (Chao convention) + if (useSpaceCharge != 0) { + double mu_0 = 4.0 * OrbitConst::PI * 1.e-07; // permeability of free space + double _z_0 = OrbitConst::c * mu_0; + + zSpaceCharge_n = _z_0 * (1.0 + 2.0 * log(b_a)) / + (2 * sp->getBeta() * sp->getGamma() * sp->getGamma()); + } + + for (int n = 1; n < nBins / 2; n++) { + realPart = std::real(_zImped_n[n]); + imagPart = std::imag(_zImped_n[n]) + zSpaceCharge_n; + _z[n] = n * sqrt(realPart * realPart + imagPart * imagPart); + _chi[n] = atan2(imagPart, realPart); + } + + // Convert charge to current for a single macroparticle per unit bin length + double charge2current = bunch->getCharge() * bunch->getMacroSize() * OrbitConst::elementary_charge_MKS * sp->getBeta() * OrbitConst::c / (length / nBins); + + // Calculate and add the kick to macroparticles + double kick, angle, position; + + // Convert particle longitudinal coordinate to phi + + double philocal; + double z; + double **coords = bunch->coordArr(); + for (int j = 0; j < bunch->getSize(); j++) { + z = bunch->z(j); + philocal = (z / length) * 2 * OrbitConst::PI; + + // Handle cases where the longitudinal coordinate is + // outside of the user-specified length + if (philocal < -OrbitConst::PI) + philocal += 2 * OrbitConst::PI; + if (philocal > OrbitConst::PI) + philocal -= 2 * OrbitConst::PI; + + double dE = _kick(philocal) * (-1e-9) * bunch->getCharge() * charge2current; + coords[j][5] += dE; + } + } + else { + std::cout << "Gradient-based solver is not implemented."; + } } - /////////////////////////////////////////////////////////////////////////// // // NAME @@ -217,20 +205,18 @@ void LSpaceChargeCalc::trackBunch(Bunch* bunch) // /////////////////////////////////////////////////////////////////////////// -double LSpaceChargeCalc::_kick(double angle) -{ - // n=0 term has no impact (constant in phi) - // f(phi) = _FFTMagnitude(1) + sum (n = 2 -> N / 2) of - // [2 * _FFTMagnitude(i) * cos(phi * (n - 1) + - // _FFTPhase(n) + _chi(n))] - - double kick = 0.; - double cosArg; - - for(int n = 1; n < nModes; n++) - { - cosArg = n * (angle + OrbitConst::PI) + _fftphase[n] + _chi[n]; - kick += 2 * _fftmagnitude[n] * _z[n] * cos(cosArg); - } - return kick; +double LSpaceChargeCalc::_kick(double angle) { + // n=0 term has no impact (constant in phi) + // f(phi) = _FFTMagnitude(1) + sum (n = 2 -> N / 2) of + // [2 * _FFTMagnitude(i) * cos(phi * (n - 1) + + // _FFTPhase(n) + _chi(n))] + + double kick = 0.; + double cosArg; + + for (int n = 1; n < nModes; n++) { + cosArg = n * (angle + OrbitConst::PI) + _fftphase[n] + _chi[n]; + kick += 2 * _fftmagnitude[n] * _z[n] * cos(cosArg); + } + return kick; } diff --git a/src/spacecharge/LSpaceChargeCalc.hh b/src/spacecharge/LSpaceChargeCalc.hh index 4e81c9e2..6f68d43e 100644 --- a/src/spacecharge/LSpaceChargeCalc.hh +++ b/src/spacecharge/LSpaceChargeCalc.hh @@ -38,7 +38,7 @@ public: void setNumModes(int n); /** Sets option to use gradient rather than impedance. **/ - void setUseGrad(int n); + void setUseGrad(int setting); /** Calculates space charge and applies the transverse and longitudinal SC kicks to the macro-particles in the bunch. */ void trackBunch(Bunch* bunch); diff --git a/src/spacecharge/wrap_lspacechargecalc.cc b/src/spacecharge/wrap_lspacechargecalc.cc index 6e04f98d..3edd68c7 100644 --- a/src/spacecharge/wrap_lspacechargecalc.cc +++ b/src/spacecharge/wrap_lspacechargecalc.cc @@ -46,8 +46,10 @@ extern "C" int nMacrosMin; int useSpaceCharge; int nBins; + int nModes; + int useGrad; - if(!PyArg_ParseTuple(args,"ddiiii:arguments", &b_a, &length, &nMacrosMin, &useSpaceCharge, &nBins)){ + if(!PyArg_ParseTuple(args,"ddiii:arguments", &b_a, &length, &nMacrosMin, &useSpaceCharge, &nBins)){ ORBIT_MPI_Finalize("PyLSpaceChargeCalc - LSpaceChargeCalc(b_a, length, nMacrosMin, useSpaceCharge, nBins) - constructor needs parameters."); } @@ -92,37 +94,31 @@ extern "C" } - static PyObject* LSpaceChargeCalc_setNumModes(PyObject *self, PyObject *args){ - + static PyObject *LSpaceChargeCalc_setNumModes(PyObject *self, PyObject *args) { pyORBIT_Object* pyLSpaceChargeCalc = (pyORBIT_Object*) self; LSpaceChargeCalc* cpp_LSpaceChargeCalc = (LSpaceChargeCalc*) pyLSpaceChargeCalc->cpp_obj; - int n = -1; - - if(!PyArg_ParseTuple(args,"i:arguments", &n)){ + int nmodes; + if (!PyArg_ParseTuple(args, "i:arguments", &nmodes)) { ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setNumModes(n) - constructor needs parameters."); - } - cpp_LSpaceChargeCalc->setNumModes(n); - - Py_INCREF(Py_None); - return Py_None; - } - - static PyObject* LSpaceChargeCalc_setUseGrad(PyObject *self, PyObject *args){ - + } + cpp_LSpaceChargeCalc->setNumModes(nmodes); + Py_INCREF(Py_None); + return Py_None; + } + + static PyObject *LSpaceChargeCalc_setUseGrad(PyObject *self, PyObject *args) { pyORBIT_Object* pyLSpaceChargeCalc = (pyORBIT_Object*) self; LSpaceChargeCalc* cpp_LSpaceChargeCalc = (LSpaceChargeCalc*) pyLSpaceChargeCalc->cpp_obj; - int setting = 0; - - if(!PyArg_ParseTuple(args,"i:arguments", &setting)){ + int usegrad; + if (!PyArg_ParseTuple(args, "i:arguments", &usegrad)) { ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setUseGrad(setting) - constructor needs parameters."); - } - cpp_LSpaceChargeCalc->setUseGrad(setting); - - Py_INCREF(Py_None); - return Py_None; - } + } + cpp_LSpaceChargeCalc->setUseGrad(usegrad); + Py_INCREF(Py_None); + return Py_None; + } //assignImpedanceValue(int, real, real). Wraps the LongSpaceChargeCalc routine assigning an impedance mode static PyObject* LSpaceChargeCalc_assignImpedanceValue(PyObject *self, PyObject *args){ @@ -180,9 +176,11 @@ extern "C" // defenition of the methods of the python LSpaceChargeCalc wrapper class // they will be vailable from python level static PyMethodDef LSpaceChargeCalcClassMethods[] = { - { "trackBunch", LSpaceChargeCalc_trackBunch, METH_VARARGS,"trackBunch the bunch - trackBunch(pyBunch)"}, - { "assignImpedanceValue", LSpaceChargeCalc_assignImpedanceValue, METH_VARARGS,"assigne the impedance for the ith mode - assignImpedanceValue(i,real,imag)"}, - { "assignImpedance", assignImpedance, METH_VARARGS,"assigne the impedance for the ith mode - assignImpedance(Z))"}, + { "trackBunch", LSpaceChargeCalc_trackBunch, METH_VARARGS,"trackBunch the bunch - trackBunch(pyBunch)"}, + { "assignImpedanceValue", LSpaceChargeCalc_assignImpedanceValue, METH_VARARGS,"assign impedance of ith mode - assignImpedanceValue(i, real, imag)"}, + { "assignImpedance", assignImpedance, METH_VARARGS, "assign impedance of ith mode - assignImpedance(Z))"}, + { "setNumModes", LSpaceChargeCalc_setNumModes, METH_VARARGS, "set number of FFT modes used to calculate energy kick"}, + { "setUseGrad", LSpaceChargeCalc_setUseGrad, METH_VARARGS, "set whether to use gradient-based solver"}, {NULL} }; From e888b4d3f94b5dc55110ae8118247c33b1244a97 Mon Sep 17 00:00:00 2001 From: 46h <46h@mac144229.ornl.gov> Date: Mon, 9 Feb 2026 16:13:59 -0500 Subject: [PATCH 09/12] Fix mpi error For some reason there is an MPI error if I call `setNumModes` when creating the `SC1D_AccNode` object. There is no error when I call this function after creating the node... --- py/orbit/space_charge/sc1d/sc1DNode.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/py/orbit/space_charge/sc1d/sc1DNode.py b/py/orbit/space_charge/sc1d/sc1DNode.py index 0bf968e1..3a751135 100644 --- a/py/orbit/space_charge/sc1d/sc1DNode.py +++ b/py/orbit/space_charge/sc1d/sc1DNode.py @@ -27,8 +27,6 @@ def __init__( nmacros_min: float, use_sc: float, nbins: float, - nmodes: int = None, - use_grad: bool = False, name="long sc node", ) -> None: """ @@ -36,8 +34,6 @@ def __init__( """ DriftTEAPOT.__init__(self, name) self.lspacecharge = LSpaceChargeCalc(b_a, phase_length, nmacros_min, use_sc, nbins) - self.setNumModes(nmodes) - # self.setUseGrad(use_grad) self.setType("long sc node") self.setLength(0.0) From 6dbacbfc6899b9198bd7eb1cac5132981d35e659 Mon Sep 17 00:00:00 2001 From: 46h <46h@mac144229.ornl.gov> Date: Mon, 9 Feb 2026 17:39:22 -0500 Subject: [PATCH 10/12] Add option to use density gradient --- src/spacecharge/LSpaceChargeCalc.cc | 31 ++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index a24f49d5..22d33231 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -184,10 +184,39 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { } } else { - std::cout << "Gradient-based solver is not implemented."; + double pi = OrbitConst::PI; + double c = OrbitConst::c; + double mu0 = OrbitConst::permeability; + double eps0 = 1.0 / (mu0 * c * c); + double g = (1.0 + 2.0 * log(b_a)); + double q = bunch->getCharge(); + + SyncPart *sp = bunch->getSyncPart(); + double beta = sp->getBeta(); + double gamma = sp->getGamma(); + double ez_factor = g * q / (4.0 * pi * eps0 * gamma * gamma); + + double kick_factor = length * bunch->getClassicalRadius() * bunch->getMass(); + + int smooth = 1; + + double z, ez, density_gradient; + for (int i = 0, n = bunch->getSize(); i < n; i++) { + z = bunch->z(i) * gamma; + if (smooth > 0) { + zGrid->calcGradientSmoothed(z, density_gradient); + } + else { + zGrid->calcGradient(z, density_gradient); + } + ez = density_gradient * ez_factor; + bunch->dE(i) -= ez * kick_factor; + } } } + + /////////////////////////////////////////////////////////////////////////// // // NAME From 26f9ec0e3476634ed9eeb9d5a40bad215092298d Mon Sep 17 00:00:00 2001 From: 46h <46h@mac144229.ornl.gov> Date: Mon, 9 Feb 2026 17:57:15 -0500 Subject: [PATCH 11/12] Add option to use smooth gradient --- py/orbit/space_charge/sc1d/sc1DNode.py | 8 ++++++-- src/spacecharge/LSpaceChargeCalc.cc | 14 ++++++++------ src/spacecharge/LSpaceChargeCalc.hh | 4 ++++ src/spacecharge/wrap_lspacechargecalc.cc | 14 ++++++++++++++ 4 files changed, 32 insertions(+), 8 deletions(-) diff --git a/py/orbit/space_charge/sc1d/sc1DNode.py b/py/orbit/space_charge/sc1d/sc1DNode.py index 3a751135..dadaae8c 100644 --- a/py/orbit/space_charge/sc1d/sc1DNode.py +++ b/py/orbit/space_charge/sc1d/sc1DNode.py @@ -37,9 +37,13 @@ def __init__( self.setType("long sc node") self.setLength(0.0) - def setUseGrad(self, use_grad: bool) -> None: + def setUseGrad(self, setting: bool) -> None: """Sets whether to use gradient-based solver instead of impedance solver.""" - self.lspacecharge.setUseGrad(int(use_grad)) + self.lspacecharge.setUseGrad(int(setting)) + + def setSmoothGrad(self, setting: bool) -> None: + """Sets whether to use smoothed gradient.""" + self.lspacecharge.setSmoothGrad(int(setting)) def setNumModes(self, n: int) -> None: """Sets number of FFT modes used to calculate energy kick.""" diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 22d33231..58d3007e 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -35,6 +35,7 @@ LSpaceChargeCalc::LSpaceChargeCalc(double b_a_in, double length_in, int nMacrosM nModes = nBins / 2; useGrad = 0; + smooth = 0; _fftmagnitude = new double[nBins / 2]; _fftphase = new double[nBins / 2]; @@ -92,6 +93,10 @@ void LSpaceChargeCalc::setUseGrad(int setting) { useGrad = setting; } +void LSpaceChargeCalc::setSmoothGrad(int setting) { + smooth = setting; +} + void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) { _zImped_n[n + 1] = std::complex(real, imag); } @@ -114,7 +119,7 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { zGrid->binBunchSmoothedByParticle(bunch); zGrid->synchronizeMPI(bunch->getMPI_Comm_Local()); - if (useGrad == 0) { + if (useGrad == 0) { // Impedance-based solver // FFT the beam density for (int i = 0; i < nBins; i++) { @@ -164,7 +169,6 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { double kick, angle, position; // Convert particle longitudinal coordinate to phi - double philocal; double z; double **coords = bunch->coordArr(); @@ -183,7 +187,7 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { coords[j][5] += dE; } } - else { + else { // Gradient-based solver double pi = OrbitConst::PI; double c = OrbitConst::c; double mu0 = OrbitConst::permeability; @@ -194,12 +198,10 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { SyncPart *sp = bunch->getSyncPart(); double beta = sp->getBeta(); double gamma = sp->getGamma(); - double ez_factor = g * q / (4.0 * pi * eps0 * gamma * gamma); + double ez_factor = g * q / (4.0 * pi * eps0 * gamma * gamma); double kick_factor = length * bunch->getClassicalRadius() * bunch->getMass(); - int smooth = 1; - double z, ez, density_gradient; for (int i = 0, n = bunch->getSize(); i < n; i++) { z = bunch->z(i) * gamma; diff --git a/src/spacecharge/LSpaceChargeCalc.hh b/src/spacecharge/LSpaceChargeCalc.hh index 6f68d43e..f9b24a70 100644 --- a/src/spacecharge/LSpaceChargeCalc.hh +++ b/src/spacecharge/LSpaceChargeCalc.hh @@ -40,6 +40,9 @@ public: /** Sets option to use gradient rather than impedance. **/ void setUseGrad(int setting); + /** Sets option to use smooth gradient of longitudinal density. **/ + void setSmoothGrad(int setting); + /** Calculates space charge and applies the transverse and longitudinal SC kicks to the macro-particles in the bunch. */ void trackBunch(Bunch* bunch); @@ -58,6 +61,7 @@ public: int useSpaceCharge; int nModes; int useGrad; + int smooth; //protected: Grid1D* zGrid; diff --git a/src/spacecharge/wrap_lspacechargecalc.cc b/src/spacecharge/wrap_lspacechargecalc.cc index 3edd68c7..fd089ad5 100644 --- a/src/spacecharge/wrap_lspacechargecalc.cc +++ b/src/spacecharge/wrap_lspacechargecalc.cc @@ -120,6 +120,19 @@ extern "C" return Py_None; } + static PyObject *LSpaceChargeCalc_setSmoothGrad(PyObject *self, PyObject *args) { + pyORBIT_Object* pyLSpaceChargeCalc = (pyORBIT_Object*) self; + LSpaceChargeCalc* cpp_LSpaceChargeCalc = (LSpaceChargeCalc*) pyLSpaceChargeCalc->cpp_obj; + + int smoothgrad; + if (!PyArg_ParseTuple(args, "i:arguments", &smoothgrad)) { + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setSmoothGrad(setting) - constructor needs parameters."); + } + cpp_LSpaceChargeCalc->setSmoothGrad(smoothgrad); + Py_INCREF(Py_None); + return Py_None; + } + //assignImpedanceValue(int, real, real). Wraps the LongSpaceChargeCalc routine assigning an impedance mode static PyObject* LSpaceChargeCalc_assignImpedanceValue(PyObject *self, PyObject *args){ @@ -181,6 +194,7 @@ extern "C" { "assignImpedance", assignImpedance, METH_VARARGS, "assign impedance of ith mode - assignImpedance(Z))"}, { "setNumModes", LSpaceChargeCalc_setNumModes, METH_VARARGS, "set number of FFT modes used to calculate energy kick"}, { "setUseGrad", LSpaceChargeCalc_setUseGrad, METH_VARARGS, "set whether to use gradient-based solver"}, + { "setSmoothGrad", LSpaceChargeCalc_setSmoothGrad, METH_VARARGS, "set whether to use smooth gradients"}, {NULL} }; From 7258f6fb3a9d097fb652b565f3ba602781cd0c13 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Fri, 13 Feb 2026 15:54:58 -0500 Subject: [PATCH 12/12] Remove if-else --- src/spacecharge/LSpaceChargeCalc.cc | 161 +++++++++++----------------- 1 file changed, 65 insertions(+), 96 deletions(-) diff --git a/src/spacecharge/LSpaceChargeCalc.cc b/src/spacecharge/LSpaceChargeCalc.cc index 58d3007e..74e40862 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -106,7 +106,7 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { if (nPartsGlobal < nMacrosMin) return; - // Bin the particles + // Bin the particles. double zmin, zmax; double realPart, imagPart; @@ -119,101 +119,70 @@ void LSpaceChargeCalc::trackBunch(Bunch *bunch) { zGrid->binBunchSmoothedByParticle(bunch); zGrid->synchronizeMPI(bunch->getMPI_Comm_Local()); - if (useGrad == 0) { // Impedance-based solver - - // FFT the beam density - for (int i = 0; i < nBins; i++) { - _in[i][0] = zGrid->getValueOnGrid(i); - _in[i][1] = 0.; - } - - fftw_execute(_plan); - - // Find the magnitude and phase - _fftmagnitude[0] = _out[0][0] / (double)nBins; - _fftphase[0] = 0.; - - for (int n = 1; n < nBins / 2; n++) { - realPart = _out[n][0] / ((double)nBins); - imagPart = _out[n][1] / ((double)nBins); - _fftmagnitude[n] = sqrt(realPart * realPart + imagPart * imagPart); - _fftphase[n] = atan2(imagPart, realPart); - } - - // Space charge contribution to impedance - SyncPart *sp = bunch->getSyncPart(); - - // Set zero if useSpaceCharge = 0 - double zSpaceCharge_n = 0.; - - // Otherwise, set positive since space charge is capacitive (Chao convention) - if (useSpaceCharge != 0) { - double mu_0 = 4.0 * OrbitConst::PI * 1.e-07; // permeability of free space - double _z_0 = OrbitConst::c * mu_0; - - zSpaceCharge_n = _z_0 * (1.0 + 2.0 * log(b_a)) / - (2 * sp->getBeta() * sp->getGamma() * sp->getGamma()); - } - - for (int n = 1; n < nBins / 2; n++) { - realPart = std::real(_zImped_n[n]); - imagPart = std::imag(_zImped_n[n]) + zSpaceCharge_n; - _z[n] = n * sqrt(realPart * realPart + imagPart * imagPart); - _chi[n] = atan2(imagPart, realPart); - } - - // Convert charge to current for a single macroparticle per unit bin length - double charge2current = bunch->getCharge() * bunch->getMacroSize() * OrbitConst::elementary_charge_MKS * sp->getBeta() * OrbitConst::c / (length / nBins); - - // Calculate and add the kick to macroparticles - double kick, angle, position; - - // Convert particle longitudinal coordinate to phi - double philocal; - double z; - double **coords = bunch->coordArr(); - for (int j = 0; j < bunch->getSize(); j++) { - z = bunch->z(j); - philocal = (z / length) * 2 * OrbitConst::PI; - - // Handle cases where the longitudinal coordinate is - // outside of the user-specified length - if (philocal < -OrbitConst::PI) - philocal += 2 * OrbitConst::PI; - if (philocal > OrbitConst::PI) - philocal -= 2 * OrbitConst::PI; - - double dE = _kick(philocal) * (-1e-9) * bunch->getCharge() * charge2current; - coords[j][5] += dE; - } - } - else { // Gradient-based solver - double pi = OrbitConst::PI; - double c = OrbitConst::c; - double mu0 = OrbitConst::permeability; - double eps0 = 1.0 / (mu0 * c * c); - double g = (1.0 + 2.0 * log(b_a)); - double q = bunch->getCharge(); - - SyncPart *sp = bunch->getSyncPart(); - double beta = sp->getBeta(); - double gamma = sp->getGamma(); - - double ez_factor = g * q / (4.0 * pi * eps0 * gamma * gamma); - double kick_factor = length * bunch->getClassicalRadius() * bunch->getMass(); - - double z, ez, density_gradient; - for (int i = 0, n = bunch->getSize(); i < n; i++) { - z = bunch->z(i) * gamma; - if (smooth > 0) { - zGrid->calcGradientSmoothed(z, density_gradient); - } - else { - zGrid->calcGradient(z, density_gradient); - } - ez = density_gradient * ez_factor; - bunch->dE(i) -= ez * kick_factor; - } + // FFT the beam density. + for (int i = 0; i < nBins; i++) { + _in[i][0] = zGrid->getValueOnGrid(i); + _in[i][1] = 0.; + } + + fftw_execute(_plan); + + // Find the magnitude and phase + _fftmagnitude[0] = _out[0][0] / (double)nBins; + _fftphase[0] = 0.; + + for (int n = 1; n < nBins / 2; n++) { + realPart = _out[n][0] / ((double)nBins); + imagPart = _out[n][1] / ((double)nBins); + _fftmagnitude[n] = sqrt(realPart * realPart + imagPart * imagPart); + _fftphase[n] = atan2(imagPart, realPart); + } + + // Space charge contribution to impedance. + SyncPart *sp = bunch->getSyncPart(); + + // Set zero if useSpaceCharge = 0. + double zSpaceCharge_n = 0.; + + // Otherwise, set positive since space charge is capacitive (Chao convention). + if (useSpaceCharge != 0) { + double mu_0 = 4.0 * OrbitConst::PI * 1.e-07; // permeability of free space + double _z_0 = OrbitConst::c * mu_0; + + zSpaceCharge_n = _z_0 * (1.0 + 2.0 * log(b_a)) / + (2 * sp->getBeta() * sp->getGamma() * sp->getGamma()); + } + + for (int n = 1; n < nBins / 2; n++) { + realPart = std::real(_zImped_n[n]); + imagPart = std::imag(_zImped_n[n]) + zSpaceCharge_n; + _z[n] = n * sqrt(realPart * realPart + imagPart * imagPart); + _chi[n] = atan2(imagPart, realPart); + } + + // Convert charge to current for a single macroparticle per unit bin length. + double charge2current = bunch->getCharge() * bunch->getMacroSize() * OrbitConst::elementary_charge_MKS * sp->getBeta() * OrbitConst::c / (length / nBins); + + // Calculate and add the kick to macroparticles. + double kick, angle, position; + + // Convert particle longitudinal coordinate to phi. + double philocal; + double z; + double **coords = bunch->coordArr(); + for (int j = 0; j < bunch->getSize(); j++) { + z = bunch->z(j); + philocal = (z / length) * 2 * OrbitConst::PI; + + // Handle cases where the longitudinal coordinate is + // outside of the user-specified length + if (philocal < -OrbitConst::PI) + philocal += 2 * OrbitConst::PI; + if (philocal > OrbitConst::PI) + philocal -= 2 * OrbitConst::PI; + + double dE = _kick(philocal) * (-1e-9) * bunch->getCharge() * charge2current; + coords[j][5] += dE; } }