diff --git a/py/orbit/space_charge/sc1d/sc1DNode.py b/py/orbit/space_charge/sc1d/sc1DNode.py index 628eaf53..dadaae8c 100644 --- a/py/orbit/space_charge/sc1d/sc1DNode.py +++ b/py/orbit/space_charge/sc1d/sc1DNode.py @@ -6,37 +6,50 @@ 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, + name="long sc node", + ) -> None: """ 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) self.setType("long sc node") self.setLength(0.0) - def trackBunch(self, bunch): + def setUseGrad(self, setting: bool) -> None: + """Sets whether to use gradient-based solver instead of impedance solver.""" + 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.""" + self.lspacecharge.setNumModes(n) + + def trackBunch(self, bunch: Bunch) -> None: """ The SC1D-teapot class implementation of the AccNodeBunchTracker class trackBunch(probe) method. @@ -44,76 +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"] - self.lspacecharge.trackBunch(bunch) # track method goes here + 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) -# ----------------------------------------------------------------------------- -# Node for impedance as function of frequency -# ----------------------------------------------------------------------------- - - class FreqDep_SC1D_AccNode(DriftTEAPOT): - def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, impeDict, name="freq. dep. long sc node"): + """Longitudinal space charge node (frequency-dependent).""" + + 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) @@ -121,23 +143,32 @@ def track(self, paramsDict): self.lspacecharge.trackBunch(bunch) -# ----------------------------------------------------------------------------- -# Node for impedance as function of beta and frequency -# ----------------------------------------------------------------------------- - - class BetFreqDep_SC1D_AccNode(DriftTEAPOT): - def __init__(self, b_a, phaseLength, nMacrosMin, useSpaceCharge, nBins, bunch, impeDict, name="freq. dep. long sc node"): + """Longitudinal space charge node (frequency- and velocity-dependent).""" + + 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"] @@ -145,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 @@ -210,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 12737969..74e40862 100644 --- a/src/spacecharge/LSpaceChargeCalc.cc +++ b/src/spacecharge/LSpaceChargeCalc.cc @@ -11,186 +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); - - _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; + smooth = 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; -void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) -{ - _zImped_n[n + 1] = std::complex(real, imag); + if (nModes > nModesMax) { + nModes = nModesMax; + } + if (nModes < 0) { + nModes = nModesMax; + } } +void LSpaceChargeCalc::setUseGrad(int setting) { + useGrad = setting; +} -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 +void LSpaceChargeCalc::setSmoothGrad(int setting) { + smooth = setting; +} - if(philocal < -OrbitConst::PI) philocal += 2 * OrbitConst::PI; - if(philocal > OrbitConst::PI) philocal -= 2 * OrbitConst::PI; +void LSpaceChargeCalc::assignImpedanceValue(int n, double real, double imag) { + _zImped_n[n + 1] = std::complex(real, imag); +} - 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()); + + // 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; + } } + /////////////////////////////////////////////////////////////////////////// // // NAME @@ -208,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 < nBins / 2; 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 24fc00a6..f9b24a70 100644 --- a/src/spacecharge/LSpaceChargeCalc.hh +++ b/src/spacecharge/LSpaceChargeCalc.hh @@ -34,12 +34,19 @@ public: /** 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 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); - /** 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,6 +59,9 @@ public: int nBins; int nMacrosMin; 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 a3a98ea2..fd089ad5 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,"ddiii: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,7 +94,44 @@ extern "C" } + static PyObject *LSpaceChargeCalc_setNumModes(PyObject *self, PyObject *args) { + pyORBIT_Object* pyLSpaceChargeCalc = (pyORBIT_Object*) self; + LSpaceChargeCalc* cpp_LSpaceChargeCalc = (LSpaceChargeCalc*) pyLSpaceChargeCalc->cpp_obj; + + int nmodes; + if (!PyArg_ParseTuple(args, "i:arguments", &nmodes)) { + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setNumModes(n) - constructor needs parameters."); + } + 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 usegrad; + if (!PyArg_ParseTuple(args, "i:arguments", &usegrad)) { + ORBIT_MPI_Finalize("PyLSpaceChargeCalc - setUseGrad(setting) - constructor needs parameters."); + } + cpp_LSpaceChargeCalc->setUseGrad(usegrad); + Py_INCREF(Py_None); + 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){ @@ -150,9 +189,12 @@ 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"}, + { "setSmoothGrad", LSpaceChargeCalc_setSmoothGrad, METH_VARARGS, "set whether to use smooth gradients"}, {NULL} };