From d95aa5cf95279090c0edd2b06f86f8f0644f04f7 Mon Sep 17 00:00:00 2001 From: lcarver Date: Mon, 4 Dec 2023 10:16:56 +0100 Subject: [PATCH] Simple Ring modification to solve energy loss computation bug (#710) * Remove U0 and damping elements from SimpleQuantDiffPass * Add simple radiation pass method * Modify SimpleQuantDiff Element and add SimpleRadiation element * Fix tests * Renamed SimpleRadPass. Modified error in lattice_object passmethod tuples, changed radiation to longdict in element, modified list of disable_6d passmethods in get_energy_loss and modified fast ring * Update len of simple_ring in test_physics which is now different * Add damping matrix * Add radiative element * Add dispersion and dispersion ' * Fix normalisation order which was problem with dispersion * remove alpha and beta, change damp_mat to keep only diagonals, add build attributes, add exponent instead of linear * Fix tests * remove alpha and beta all over, but keeping dispersion * bug in mexall * Add dispersion to linear matrix * Add optional for disp --------- Co-authored-by: Lee Carver Co-authored-by: Lee Carver --- atintegrators/SimpleQuantDiffPass.c | 53 ++-------- atintegrators/SimpleRadiationPass.c | 153 ++++++++++++++++++++++++++++ pyat/at/lattice/elements.py | 59 ++++++++++- pyat/at/lattice/lattice_object.py | 6 +- pyat/at/physics/energy_loss.py | 4 +- pyat/at/physics/fastring.py | 40 ++++++-- pyat/test/test_physics.py | 2 +- 7 files changed, 251 insertions(+), 66 deletions(-) create mode 100644 atintegrators/SimpleRadiationPass.c diff --git a/atintegrators/SimpleQuantDiffPass.c b/atintegrators/SimpleQuantDiffPass.c index 2bc5895dd..988fad0be 100644 --- a/atintegrators/SimpleQuantDiffPass.c +++ b/atintegrators/SimpleQuantDiffPass.c @@ -14,13 +14,11 @@ struct elem double betay; double sigma_xp; double sigma_yp; - double U0; - double EnergyLossFactor; }; void SimpleQuantDiffPass(double *r_in, double sigma_xp, double sigma_yp, double espread, - double taux, double tauy, double tauz, double EnergyLossFactor, + double taux, double tauy, double tauz, pcg32_random_t* rng, int num_particles) { @@ -36,28 +34,15 @@ void SimpleQuantDiffPass(double *r_in, } if(!atIsNaN(r6[0])) { - if(taux!=0.0) { - r6[1] -= 2*r6[1]/taux; - } if(sigma_xp!=0.0) { r6[1] += 2*sigma_xp*sqrt(1/taux)*randnorm[0]; } if(sigma_yp!=0.0) { r6[3] += 2*sigma_yp*sqrt(1/tauy)*randnorm[1]; } - if(tauy!=0.0) { - r6[3] -= 2*r6[3]/tauy; - } if(espread!=0.0) { r6[4] += 2*espread*sqrt(1/tauz)*randnorm[2]; } - if(tauz!=0.0) { - r6[4] -= 2*r6[4]/tauz; - } - - if(EnergyLossFactor>0.0) { - r6[4] -= EnergyLossFactor; - } } } } @@ -68,7 +53,7 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem, { /* if (ElemData) {*/ if (!Elem) { - double emitx, emity, espread, taux, tauy, tauz, betax, betay, U0, EnergyLossFactor; + double emitx, emity, espread, taux, tauy, tauz, betax, betay; emitx=atGetDouble(ElemData,"emitx"); check_error(); emity=atGetDouble(ElemData,"emity"); check_error(); espread=atGetDouble(ElemData,"espread"); check_error(); @@ -77,7 +62,6 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem, tauz=atGetDouble(ElemData,"tauz"); check_error(); betax=atGetDouble(ElemData,"betax"); check_error(); betay=atGetDouble(ElemData,"betay"); check_error(); - U0=atGetDouble(ElemData,"U0"); check_error(); Elem = (struct elem*)atMalloc(sizeof(struct elem)); Elem->emitx=emitx; @@ -90,19 +74,8 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem, Elem->betay=betay; Elem->sigma_xp=sqrt(emitx/betax); Elem->sigma_yp=sqrt(emity/betay); - Elem->U0=U0; - Elem->EnergyLossFactor=U0/Param->energy; } - SimpleQuantDiffPass(r_in, Elem->sigma_xp, Elem->sigma_yp, Elem->espread, Elem->taux, Elem->tauy, Elem->tauz, Elem->EnergyLossFactor, Param->thread_rng, num_particles); -/* } - else { - atFree(Elem->T1); - atFree(Elem->T2); - atFree(Elem->R1); - atFree(Elem->R2); - atFree(Elem->EApertures); - atFree(Elem->RApertures); - }*/ + SimpleQuantDiffPass(r_in, Elem->sigma_xp, Elem->sigma_yp, Elem->espread, Elem->taux, Elem->tauy, Elem->tauz, Param->thread_rng, num_particles); return Elem; } @@ -117,7 +90,7 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) double *r_in; const mxArray *ElemData = prhs[0]; int num_particles = mxGetN(prhs[1]); - double emitx, emity, espread, taux, tauy, tauz, betax, betay, sigma_xp, sigma_yp, U0, EnergyLossFactor; + double emitx, emity, espread, taux, tauy, tauz, betax, betay, sigma_xp, sigma_yp; emitx=atGetDouble(ElemData,"emitx"); check_error(); emity=atGetDouble(ElemData,"emity"); check_error(); @@ -127,19 +100,17 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) tauz=atGetDouble(ElemData,"tauz"); check_error(); betax=atGetDouble(ElemData,"betax"); check_error(); betay=atGetDouble(ElemData,"betay"); check_error(); - U0=atGetDouble(ElemData,"U0"); check_error(); sigma_xp=sqrt(emitx/betax); sigma_yp=sqrt(emity/betay); - EnergyLossFactor=U0/6e9; if (mxGetM(prhs[1]) != 6) mexErrMsgIdAndTxt("AT:WrongArg","Second argument must be a 6 x N matrix"); /* ALLOCATE memory for the output array of the same size as the input */ plhs[0] = mxDuplicateArray(prhs[1]); r_in = mxGetDoubles(plhs[0]); - SimpleQuantDiffPass(r_in, sigma_xp, sigma_yp, espread, taux, tauy, tauz, EnergyLossFactor, &pcg32_global, num_particles); + SimpleQuantDiffPass(r_in, sigma_xp, sigma_yp, espread, taux, tauy, tauz, &pcg32_global, num_particles); } else if (nrhs == 0) { /* list of required fields */ - plhs[0] = mxCreateCellMatrix(9,1); + plhs[0] = mxCreateCellMatrix(8,1); mxSetCell(plhs[0],0,mxCreateString("emitx")); mxSetCell(plhs[0],1,mxCreateString("emity")); mxSetCell(plhs[0],2,mxCreateString("espread")); @@ -148,19 +119,7 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) mxSetCell(plhs[0],5,mxCreateString("tauz")); mxSetCell(plhs[0],6,mxCreateString("betax")); mxSetCell(plhs[0],7,mxCreateString("betay")); - mxSetCell(plhs[0],8,mxCreateString("U0")); - - /* - if (nlhs>1) { - plhs[1] = mxCreateCellMatrix(6,1); - mxSetCell(plhs[1],0,mxCreateString("T1")); - mxSetCell(plhs[1],1,mxCreateString("T2")); - mxSetCell(plhs[1],2,mxCreateString("R1")); - mxSetCell(plhs[1],3,mxCreateString("R2")); - mxSetCell(plhs[1],4,mxCreateString("RApertures")); - mxSetCell(plhs[1],5,mxCreateString("EApertures")); - }*/ } else { mexErrMsgIdAndTxt("AT:WrongArg","Needs 0 or 2 arguments"); diff --git a/atintegrators/SimpleRadiationPass.c b/atintegrators/SimpleRadiationPass.c new file mode 100644 index 000000000..2bed3b515 --- /dev/null +++ b/atintegrators/SimpleRadiationPass.c @@ -0,0 +1,153 @@ + +#include "atelem.c" +#include "atrandom.c" + +struct elem +{ + double *damp_mat_diag; + double dispx; + double dispxp; + double dispy; + double dispyp; + double U0; + double EnergyLossFactor; +}; + +void SimpleRadiationPass(double *r_in, + double *damp_mat_diag, double dispx, double dispxp, + double dispy, double dispyp, double EnergyLossFactor, int num_particles) + +{ + double *r6; + int c, i; + double x,xp,y,yp,z,dpp; + + #pragma omp parallel for if (num_particles > OMP_PARTICLE_THRESHOLD*10) default(shared) shared(r_in,num_particles) private(c,r6) + for (c = 0; cU0=U0; + Elem->dispx=dispx; + Elem->dispxp=dispxp; + Elem->dispy=dispy; + Elem->dispyp=dispyp; + Elem->EnergyLossFactor=U0/Param->energy; + Elem->damp_mat_diag=damp_mat_diag; + } + SimpleRadiationPass(r_in, Elem->damp_mat_diag, Elem->dispx, Elem->dispxp, Elem->dispy, Elem->dispyp, Elem->EnergyLossFactor, num_particles); + return Elem; +} + +MODULE_DEF(SimpleRadiationPass) /* Dummy module initialisation */ + +#endif /*defined(MATLAB_MEX_FILE) || defined(PYAT)*/ + +#if defined(MATLAB_MEX_FILE) +void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) +{ + if (nrhs == 2) { + double *r_in; + const mxArray *ElemData = prhs[0]; + int num_particles = mxGetN(prhs[1]); + double U0, EnergyLossFactor; + double dispx, dispxp, dispy, dispyp; + double *damp_mat_diag; + + damp_mat_diag=atGetDoubleArray(ElemData,"damp_mat_diag"); check_error(); + dispx=atGetOptionalDouble(ElemData,"dispx",0.0); check_error(); + dispy=atGetOptionalDouble(ElemData,"dispy",0.0); check_error(); + dispxp=atGetOptionalDouble(ElemData,"dispxp",0.0); check_error(); + dispyp=atGetOptionalDouble(ElemData,"dispyp",0.0); check_error(); + U0=atGetDouble(ElemData,"U0"); check_error(); + EnergyLossFactor=U0/6e9; + if (mxGetM(prhs[1]) != 6) mexErrMsgIdAndTxt("AT:WrongArg","Second argument must be a 6 x N matrix"); + /* ALLOCATE memory for the output array of the same size as the input */ + plhs[0] = mxDuplicateArray(prhs[1]); + r_in = mxGetDoubles(plhs[0]); + SimpleRadiationPass(r_in, damp_mat_diag, dispx, dispxp, dispy, dispyp, EnergyLossFactor, num_particles); + } + else if (nrhs == 0) { + /* list of required fields */ + plhs[0] = mxCreateCellMatrix(2,1); + mxSetCell(plhs[0],0,mxCreateString("damp_mat_diag")); + mxSetCell(plhs[0],1,mxCreateString("U0")); + + if (nlhs>1) { + /* list of optional fields */ + plhs[1] = mxCreateCellMatrix(4,1); /* No optional fields */ + mxSetCell(plhs[1],0,mxCreateString("dispx")); + mxSetCell(plhs[1],1,mxCreateString("dispxp")); + mxSetCell(plhs[1],2,mxCreateString("dispy")); + mxSetCell(plhs[1],3,mxCreateString("dispyp")); + } + } + else { + mexErrMsgIdAndTxt("AT:WrongArg","Needs 0 or 2 arguments"); + } +} +#endif /*defined(MATLAB_MEX_FILE)*/ diff --git a/pyat/at/lattice/elements.py b/pyat/at/lattice/elements.py index 905d9916d..342c0e0cc 100644 --- a/pyat/at/lattice/elements.py +++ b/pyat/at/lattice/elements.py @@ -1102,7 +1102,10 @@ def __init__(self, family_name: str, m66=None, **kwargs): class SimpleQuantDiff(_DictLongtMotion, Element): """ Linear tracking element for a simplified quantum diffusion, - radiation damping and energy loss + radiation damping and energy loss. + + Note: The damping times are needed to compute the correct + kick for the emittance. Radiation damping is NOT applied. """ _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES default_pass = {False: 'IdentityPass', True: 'SimpleQuantDiffPass'} @@ -1111,7 +1114,7 @@ def __init__(self, family_name: str, betax: float = 1.0, betay: float = 1.0, emitx: float = 0.0, emity: float = 0.0, espread: float = 0.0, taux: float = 0.0, tauy: float = 0.0, - tauz: float = 0.0, U0: float = 0.0, + tauz: float = 0.0, **kwargs): """ Args: @@ -1126,7 +1129,6 @@ def __init__(self, family_name: str, betax: float = 1.0, taux: Horizontal damping time [turns] tauy: Vertical damping time [turns] tauz: Longitudinal damping time [turns] - U0: Energy Loss [eV] Default PassMethod: ``SimpleQuantDiffPass`` """ @@ -1156,12 +1158,61 @@ def __init__(self, family_name: str, betax: float = 1.0, if espread > 0.0: assert tauz > 0.0, 'if espread is given, tauz must be non zero' - self.U0 = U0 self.betax = betax self.betay = betay super(SimpleQuantDiff, self).__init__(family_name, **kwargs) +class SimpleRadiation(_DictLongtMotion, Radiative, Element): + """Simple radiation damping and energy loss""" + _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ['taux', 'tauy', 'tauz'] + default_pass = {False: 'IdentityPass', True: 'SimpleRadiationPass'} + + def __init__(self, family_name: str, + taux: float = 0.0, tauy: float = 0.0, + tauz: float = 0.0, U0: float = 0.0, + **kwargs): + """ + Args: + family_name: Name of the element + + Optional Args: + taux: Horizontal damping time [turns] + tauy: Vertical damping time [turns] + tauz: Longitudinal damping time [turns] + U0: Energy loss per turn [eV] + + Default PassMethod: ``SimpleRadiationPass`` + """ + kwargs.setdefault('PassMethod', self.default_pass[True]) + + assert taux >= 0.0, 'taux must be greater than or equal to 0' + if taux == 0.0: + dampx = 1 + else: + dampx = numpy.exp(-1/taux) + + assert tauy >= 0.0, 'tauy must be greater than or equal to 0' + if tauy == 0.0: + dampy = 1 + else: + dampy = numpy.exp(-1/tauy) + + assert tauz >= 0.0, 'tauz must be greater than or equal to 0' + if tauz == 0.0: + dampz = 1 + else: + dampz = numpy.exp(-1/tauz) + + self.U0 = U0 + + self.damp_mat_diag = numpy.array([dampx, dampx, + dampy, dampy, + dampz, dampz]) + + super(SimpleRadiation, self).__init__(family_name, **kwargs) + + class Corrector(LongElement): """Corrector element""" _BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['KickAngle'] diff --git a/pyat/at/lattice/lattice_object.py b/pyat/at/lattice/lattice_object.py index e122f4069..a9972f967 100644 --- a/pyat/at/lattice/lattice_object.py +++ b/pyat/at/lattice/lattice_object.py @@ -50,7 +50,8 @@ ('collective_pass', elt.Collective, 'auto'), ('diffusion_pass', elt.QuantumDiffusion, 'auto'), ('energyloss_pass', elt.EnergyLoss, 'auto'), - ('simplequantdiff_pass', elt.SimpleQuantDiff, 'auto') + ('simplequantdiff_pass', elt.SimpleQuantDiff, 'auto'), + ('simpleradiation_pass', elt.SimpleRadiation, 'auto') ), True: ( ('cavity_pass', elt.RFCavity, 'auto'), @@ -63,7 +64,8 @@ ('collective_pass', elt.Collective, 'auto'), ('diffusion_pass', elt.QuantumDiffusion, 'auto'), ('energyloss_pass', elt.EnergyLoss, 'auto'), - ('simplequantdiff_pass', elt.SimpleQuantDiff, 'auto') + ('simplequantdiff_pass', elt.SimpleQuantDiff, 'auto'), + ('simpleradiation_pass', elt.SimpleRadiation, 'auto') ) } diff --git a/pyat/at/physics/energy_loss.py b/pyat/at/physics/energy_loss.py index 52b1260ac..f377b8568 100644 --- a/pyat/at/physics/energy_loss.py +++ b/pyat/at/physics/energy_loss.py @@ -6,7 +6,7 @@ from scipy.optimize import least_squares from at.lattice import Lattice, Dipole, Wiggler, RFCavity, Refpts, EnergyLoss from at.lattice import check_radiation, AtError, AtWarning -from at.lattice import QuantumDiffusion, Collective +from at.lattice import QuantumDiffusion, Collective, SimpleQuantDiff from at.lattice import get_bool_index, set_value_refpts from at.constants import clight, Cgamma from at.tracking import internal_lpass @@ -76,7 +76,7 @@ def tracking(ring): """Losses from tracking """ ringtmp = ring.disable_6d(RFCavity, QuantumDiffusion, Collective, - copy=True) + SimpleQuantDiff, copy=True) o6 = numpy.squeeze(internal_lpass(ringtmp, numpy.zeros(6), refpts=len(ringtmp))) if numpy.isnan(o6[0]): diff --git a/pyat/at/physics/fastring.py b/pyat/at/physics/fastring.py index e45f2c93d..c4be904f5 100644 --- a/pyat/at/physics/fastring.py +++ b/pyat/at/physics/fastring.py @@ -5,7 +5,7 @@ import numpy from typing import Tuple, Optional from at.lattice import RFCavity, Element, Marker, Lattice, get_cells, checkname -from at.lattice import get_elements, M66, SimpleQuantDiff, AtError +from at.lattice import get_elements, M66, SimpleQuantDiff, AtError, SimpleRadiation from at.physics import gen_m66_elem, gen_detuning_elem, gen_quantdiff_elem from at.constants import clight, e_mass import copy @@ -107,6 +107,8 @@ def simple_ring(energy: float, circumference: float, harmonic_number: int, Qx: float, Qy: float, Vrf: float, alpha: float, betax: float = 1.0, betay: float = 1.0, alphax: float = 0.0, alphay: float = 0.0, + dispx: float = 0.0, dispxp: float = 0.0, + dispy: float = 0.0, dispyp: float = 0.0, Qpx: float = 0.0, Qpy: float = 0.0, A1: float = 0.0, A2: float = 0.0, A3: float = 0.0, emitx: float = 0.0, @@ -121,10 +123,11 @@ def simple_ring(energy: float, circumference: float, harmonic_number: int, A simple ring consists of: * an RF cavity, - * a 6x6 linear transfer map, + * a 6x6 linear transfer map with no radiation damping, + * a simple radiation damping element * a detuning and chromaticity element, * a simplified quantum diffusion element - which contains equilibrium emittance and radiation damping + which contains equilibrium emittance Positional Arguments: * energy [eV] @@ -141,6 +144,10 @@ def simple_ring(energy: float, circumference: float, harmonic_number: int, * betay: vertical beta function [m], Default=1 * alphax: horizontal alpha function, Default=0 * alphay: vertical alpha function, Default=0 + * dispx: horizontal dispersion [m], Default=0 + * dispxp: horizontal dispersion prime, Default=0 + * dispy: vertical dispersion [m], Default=0 + * dispyp: vertical dispersion prime, Default=0 * Qpx: horizontal linear chromaticity, Default=0 * Qpy: vertical linear chromaticity, Default=0 * A1: horizontal amplitude detuning coefficient, Default=0 @@ -206,33 +213,46 @@ def simple_ring(energy: float, circumference: float, harmonic_number: int, M10 = -(1. + alphax**2) / betax * s_dphi_x M11 = c_dphi_x - alphax * s_dphi_x + M04 = (1 - M00) * dispx - M01 * dispxp + M14 = -M10 * dispx + (1 - M11) * dispxp + M22 = c_dphi_y + alphay * s_dphi_y M23 = betay * s_dphi_y M32 = -(1. + alphay**2) / betay * s_dphi_y M33 = c_dphi_y - alphay * s_dphi_y + M24 = (1 - M22) * dispy - M23 * dispyp + M34 = -M32 * dispy + (1 - M33) * dispyp + + M44 = 1. M45 = 0. M54 = eta*circumference M55 = 1 - Mat66 = numpy.array([[M00, M01, 0., 0., 0., 0.], - [M10, M11, 0., 0., 0., 0.], - [0., 0., M22, M23, 0., 0.], - [0., 0., M32, M33, 0., 0.], + Mat66 = numpy.array([[M00, M01, 0., 0., M04, 0.], + [M10, M11, 0., 0., M14, 0.], + [0., 0., M22, M23, M24, 0.], + [0., 0., M32, M33, M34, 0.], [0., 0., 0., 0., M44, M45], [0., 0., 0., 0., M54, M55]], order='F') # generate the linear tracking element, we set a length # which is needed to give the lattice object the correct length - # (although it is not used) + # (although it is not used for anything else) lin_elem = M66('Linear', m66=Mat66, Length=circumference) + # Generate the simple radiation element + simplerad = SimpleRadiation('SR', taux=taux, tauy=tauy, + tauz=tauz, U0=U0, dispx=dispx, + dispy=dispy, dispxp=dispxp, + dispyp=dispyp) + # Generate the simple quantum diffusion element quantdiff = SimpleQuantDiff('SQD', betax=betax, betay=betay, emitx=emitx, emity=emity, espread=espread, taux=taux, - tauy=tauy, tauz=tauz, U0=U0) + tauy=tauy, tauz=tauz) # Generate the detuning element nonlin_elem = Element('NonLinear', PassMethod='DeltaQPass', @@ -242,7 +262,7 @@ def simple_ring(energy: float, circumference: float, harmonic_number: int, A1=A1, A2=A2, A3=A3) # Assemble all elements into the lattice object - ring = Lattice(all_cavities + [lin_elem, nonlin_elem, quantdiff], + ring = Lattice(all_cavities + [lin_elem, nonlin_elem, simplerad, quantdiff], energy=energy, periodicity=1) return ring diff --git a/pyat/test/test_physics.py b/pyat/test/test_physics.py index e32bf3703..8d98710d3 100644 --- a/pyat/test/test_physics.py +++ b/pyat/test/test_physics.py @@ -325,7 +325,7 @@ def test_quantdiff(hmba_lattice): def test_simple_ring(): ring = physics.simple_ring(6e9, 844, 992, 0.1, 0.2, 6e6, 8.5e-5) - assert_equal(len(ring), 4) + assert_equal(len(ring), 5) assert_equal(ring[-1].PassMethod, 'SimpleQuantDiffPass') ring.disable_6d() assert_equal(ring[-1].PassMethod, 'IdentityPass')