From 8b8fca8cdd334637f5e28eb7a2783197b481b9ea Mon Sep 17 00:00:00 2001 From: lcarver Date: Sun, 27 Aug 2023 16:42:08 +0200 Subject: [PATCH] Simple ring model (#643) * New PassMethod for simple quantum diffusion. Added element to elements.py. Added function in fastring to generate ring * Wrong function name in __all__ * fix matlab test * matlab test * matlab tests * Add radiation damping and U0 to SimpleQuantDiffPass. * Update arguments for SimpleQuantDiff Element * Remove energyloss element. Remove analytical tau. * Add rad on/off switching. Modify input types. Set defaults * pep * change default passmethod of simplequantdiff * Add periodicity=1 to avoid at warning about no bends * Change type hint for harmonic_number to int instead of float * expanded help * wrong text * Modified text * Added multiple cavity inputs * Remved sincos * Add test. Modified help of simple_ring. Added broadcasting and optional TimeLag for cavity settings --------- Co-authored-by: Lee Carver Co-authored-by: Lee Carver --- atintegrators/SimpleQuantDiffPass.c | 169 ++++++++++++++++++++++++++++ pyat/at/lattice/elements.py | 79 +++++++++++-- pyat/at/lattice/lattice_object.py | 2 + pyat/at/physics/fastring.py | 157 +++++++++++++++++++++++++- pyat/test/test_physics.py | 10 +- 5 files changed, 403 insertions(+), 14 deletions(-) create mode 100644 atintegrators/SimpleQuantDiffPass.c diff --git a/atintegrators/SimpleQuantDiffPass.c b/atintegrators/SimpleQuantDiffPass.c new file mode 100644 index 000000000..a457aef96 --- /dev/null +++ b/atintegrators/SimpleQuantDiffPass.c @@ -0,0 +1,169 @@ + +#include "atelem.c" +#include "atrandom.c" + +struct elem +{ + double emit_x; + double emit_y; + double sigma_dp; + double tau_x; + double tau_y; + double tau_z; + double beta_x; + double beta_y; + double sigma_xp; + double sigma_yp; + double U0; + double EnergyLossFactor; +}; + +void SimpleQuantDiffPass(double *r_in, + double sigma_xp, double sigma_yp, double sigma_dp, + double tau_x, double tau_y, double tau_z, double EnergyLossFactor, + pcg32_random_t* rng, int num_particles) + +{ + double *r6; + int c, i; + + #pragma omp parallel for if (num_particles > OMP_PARTICLE_THRESHOLD*10) default(shared) shared(r_in,num_particles) private(c,r6) + for (c = 0; c0.0) { + r6[4] -= EnergyLossFactor; + } + } + } +} + +#if defined(MATLAB_MEX_FILE) || defined(PYAT) +ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem, + double *r_in, int num_particles, struct parameters *Param) +{ +/* if (ElemData) {*/ + if (!Elem) { + double emit_x, emit_y, sigma_dp, tau_x, tau_y, tau_z, beta_x, beta_y, U0, EnergyLossFactor; + emit_x=atGetDouble(ElemData,"emit_x"); check_error(); + emit_y=atGetDouble(ElemData,"emit_y"); check_error(); + sigma_dp=atGetDouble(ElemData,"sigma_dp"); check_error(); + tau_x=atGetDouble(ElemData,"tau_x"); check_error(); + tau_y=atGetDouble(ElemData,"tau_y"); check_error(); + tau_z=atGetDouble(ElemData,"tau_z"); check_error(); + beta_x=atGetDouble(ElemData,"beta_x"); check_error(); + beta_y=atGetDouble(ElemData,"beta_y"); check_error(); + U0=atGetDouble(ElemData,"U0"); check_error(); + + Elem = (struct elem*)atMalloc(sizeof(struct elem)); + Elem->emit_x=emit_x; + Elem->emit_y=emit_y; + Elem->sigma_dp=sigma_dp; + Elem->tau_x=tau_x; + Elem->tau_y=tau_y; + Elem->tau_z=tau_z; + Elem->beta_x=beta_x; + Elem->beta_y=beta_y; + Elem->sigma_xp=sqrt(emit_x/beta_x); + Elem->sigma_yp=sqrt(emit_y/beta_y); + Elem->U0=U0; + Elem->EnergyLossFactor=U0/Param->energy; + } + SimpleQuantDiffPass(r_in, Elem->sigma_xp, Elem->sigma_yp, Elem->sigma_dp, Elem->tau_x, Elem->tau_y, Elem->tau_z, 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); + }*/ + return Elem; +} + +MODULE_DEF(SimpleQuantDiffPass) /* 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 emit_x, emit_y, sigma_dp, tau_x, tau_y, tau_z, beta_x, beta_y, sigma_xp, sigma_yp, U0, EnergyLossFactor; + + emit_x=atGetDouble(ElemData,"emit_x"); check_error(); + emit_y=atGetDouble(ElemData,"emit_y"); check_error(); + sigma_dp=atGetDouble(ElemData,"sigma_dp"); check_error(); + tau_x=atGetDouble(ElemData,"tau_x"); check_error(); + tau_y=atGetDouble(ElemData,"tau_y"); check_error(); + tau_z=atGetDouble(ElemData,"tau_z"); check_error(); + beta_x=atGetDouble(ElemData,"beta_x"); check_error(); + beta_y=atGetDouble(ElemData,"beta_y"); check_error(); + U0=atGetDouble(ElemData,"U0"); check_error(); + sigma_xp=sqrt(emit_x/beta_x); + sigma_yp=sqrt(emit_y/beta_y); + 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, sigma_dp, tau_x, tau_y, tau_z, EnergyLossFactor, &pcg32_global, num_particles); + } + else if (nrhs == 0) { + /* list of required fields */ + plhs[0] = mxCreateCellMatrix(9,1); + mxSetCell(plhs[0],0,mxCreateString("emit_x")); + mxSetCell(plhs[0],1,mxCreateString("emit_y")); + mxSetCell(plhs[0],2,mxCreateString("sigma_dp")); + mxSetCell(plhs[0],3,mxCreateString("tau_x")); + mxSetCell(plhs[0],4,mxCreateString("tau_y")); + mxSetCell(plhs[0],5,mxCreateString("tau_z")); + mxSetCell(plhs[0],6,mxCreateString("beta_x")); + mxSetCell(plhs[0],7,mxCreateString("beta_y")); + 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"); + } +} +#endif /*defined(MATLAB_MEX_FILE)*/ diff --git a/pyat/at/lattice/elements.py b/pyat/at/lattice/elements.py index b0e238937..c018b3243 100644 --- a/pyat/at/lattice/elements.py +++ b/pyat/at/lattice/elements.py @@ -205,8 +205,8 @@ class Radiative(_Radiative): class Collective(_DictLongtMotion): """Mixin class for elements representing collective effects - Derived classes will automatically set the :py:attr:`~Element.is_collective` - property when the element is active. + Derived classes will automatically set the + :py:attr:`~Element.is_collective` property when the element is active. The class must have a ``default_pass`` class attribute, a dictionary such that: @@ -230,7 +230,7 @@ class Collective(_DictLongtMotion): def _get_collective(self): # noinspection PyUnresolvedReferences return self.PassMethod != self.default_pass[False] - + @abc.abstractmethod def clear_history(self): pass @@ -330,7 +330,7 @@ def swapattr(element, attro, attri): val = getattr(element, attri) delattr(element, attri) return attro, val - if copy: + if copy: el = self.copy() else: el = self @@ -482,14 +482,14 @@ def __init__(self, family_name: str, **kwargs): def set_buffers(self, nturns, nbunch): self._stds = numpy.zeros((6, nbunch, nturns), order='F') self._means = numpy.zeros((6, nbunch, nturns), order='F') - + @property def stds(self): return self._stds - + @property def means(self): - return self._means + return self._means class Aperture(Element): @@ -983,6 +983,70 @@ def __init__(self, family_name: str, m66=None, **kwargs): super(M66, self).__init__(family_name, M66=m66, **kwargs) +class SimpleQuantDiff(_DictLongtMotion, Element): + """ + Linear tracking element for a simplified quantum diffusion, + radiation damping and energy loss + """ + _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + default_pass = {False: 'IdentityPass', True: 'SimpleQuantDiffPass'} + + def __init__(self, family_name: str, beta_x: Optional[float]=1.0, + beta_y: Optional[float]=1.0, emit_x: Optional[float]=0.0, + emit_y: Optional[float]=0.0, sigma_dp: Optional[float]=0.0, + tau_x: Optional[float]=0.0, tau_y: Optional[float]=0.0, + tau_z: Optional[float]=0.0, U0: Optional[float]=0.0, + **kwargs): + """ + Args: + family_name: Name of the element + + Optional Args: + beta_x: Horizontal beta function at element [m] + beta_y: Vertical beta function at element [m] + emit_x: Horizontal equilibrium emittance [m.rad] + emit_y: Vertical equilibrium emittance [m.rad] + sigma_dp: Equilibrium energy spread + tau_x: Horizontal damping time [turns] + tau_y: Vertical damping time [turns] + tau_z: Longitudinal damping time [turns] + U0: Energy Loss [eV] + + Default PassMethod: ``SimpleQuantDiffPass`` + """ + kwargs.setdefault('PassMethod', self.default_pass[True]) + + assert tau_x>=0.0, 'tau_x must be greater than or equal to 0' + self.tau_x = tau_x + + assert tau_y>=0.0, 'tau_y must be greater than or equal to 0' + self.tau_y = tau_y + + assert tau_z>=0.0, 'tau_z must be greater than or equal to 0' + self.tau_z = tau_z + + assert emit_x>=0.0, 'emit_x must be greater than or equal to 0' + self.emit_x = emit_x + if emit_x>0.0: + assert tau_x>0.0, 'if emit_x is given, tau_x must be non zero' + + assert emit_y>=0.0, 'emit_x must be greater than or equal to 0' + self.emit_y = emit_y + if emit_y>0.0: + assert tau_y>0.0, 'if emit_y is given, tau_y must be non zero' + + assert sigma_dp>=0.0, 'sigma_dp must be greater than or equal to 0' + self.sigma_dp = sigma_dp + if sigma_dp>0.0: + assert tau_z>0.0, 'if sigma_dp is given, tau_z must be non zero' + + self.U0 = U0 + self.beta_x = beta_x + self.beta_y = beta_y + super(SimpleQuantDiff, self).__init__(family_name, **kwargs) + + + class Corrector(LongElement): """Corrector element""" _BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['KickAngle'] @@ -1100,7 +1164,6 @@ def __init__(self, family_name: str, energy_loss: float, **kwargs): kwargs.setdefault('PassMethod', self.default_pass[False]) super().__init__(family_name, EnergyLoss=energy_loss, **kwargs) - Radiative.register(EnergyLoss) diff --git a/pyat/at/lattice/lattice_object.py b/pyat/at/lattice/lattice_object.py index e9eaa417f..e122f4069 100644 --- a/pyat/at/lattice/lattice_object.py +++ b/pyat/at/lattice/lattice_object.py @@ -50,6 +50,7 @@ ('collective_pass', elt.Collective, 'auto'), ('diffusion_pass', elt.QuantumDiffusion, 'auto'), ('energyloss_pass', elt.EnergyLoss, 'auto'), + ('simplequantdiff_pass', elt.SimpleQuantDiff, 'auto') ), True: ( ('cavity_pass', elt.RFCavity, 'auto'), @@ -62,6 +63,7 @@ ('collective_pass', elt.Collective, 'auto'), ('diffusion_pass', elt.QuantumDiffusion, 'auto'), ('energyloss_pass', elt.EnergyLoss, 'auto'), + ('simplequantdiff_pass', elt.SimpleQuantDiff, 'auto') ) } diff --git a/pyat/at/physics/fastring.py b/pyat/at/physics/fastring.py index 7e5c9af42..72859fe16 100644 --- a/pyat/at/physics/fastring.py +++ b/pyat/at/physics/fastring.py @@ -3,14 +3,15 @@ """ from functools import reduce import numpy -from typing import Tuple -from at.lattice import RFCavity, Marker, Lattice, get_cells, checkname -from at.lattice import get_elements +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.physics import gen_m66_elem, gen_detuning_elem, gen_quantdiff_elem +from at.constants import clight, e_mass import copy -__all__ = ['fast_ring'] +__all__ = ['fast_ring', 'simple_ring'] def _rearrange(ring: Lattice, split_inds=[]): @@ -63,7 +64,7 @@ def _fring(ring, split_inds=[], detuning_elem=None): try: qd_elem = gen_quantdiff_elem(merged_ring) fastring.append(qd_elem) - except ValueError: # No synchrotron radiation => no diffusion element + except ValueError: # No synchrotron radiation => no diffusion element pass fastring = Lattice(fastring, **vars(ring)) return fastring @@ -100,3 +101,149 @@ def fast_ring(ring: Lattice, split_inds=[]) -> Tuple[Lattice, Lattice]: split_inds=split_inds, detuning_elem=detuning_elem) return fastringnorad, fastringrad + + +def simple_ring(energy: float, circumference: float, harmonic_number: int, + Qx: float, Qy: float, Vrf: float, alpha: float, + beta_x: Optional[float]=1.0, beta_y: Optional[float]=1.0, + alpha_x: Optional[float]=0.0, alpha_y: Optional[float]=0.0, + Qpx: Optional[float]=0.0, Qpy: Optional[float]=0.0, + A1: Optional[float]=0.0, A2: Optional[float]=0.0, + A3: Optional[float]=0.0, emit_x: Optional[float]=0.0, + emit_y: Optional[float]=0.0, sigma_dp: Optional[float]=0.0, + tau_x: Optional[float]=0.0, tau_y: Optional[float]=0.0, + tau_z: Optional[float]=0.0, U0: Optional[float]=0.0, + TimeLag: Optional[bool]=False + ): + """Generates a "simple ring" based on a given dictionary + of global parameters + + A simple ring consists of: + + * an RF cavity, + * a 6x6 linear transfer map, + * a detuning and chromaticity element, + * a simplified quantum diffusion element + which contains equilibrium emittance and radiation damping + + Positional Arguments: + * energy [eV] + * circumference [m] + * harmonic_number - can be scalar or sequence of scalars. The RF + frequency is derived from this and the ring circumference + * Qx - horizontal tune + * Qy - vertical tune + * Vrf - RF Voltage set point [V] - can be scalar or sequence of scalars + * alpha (momentum compaction factor) + + Optional Arguments: + * beta_x: horizontal beta function [m], Default=1 + * beta_y: vertical beta function [m], Default=1 + * alpha_x: horizontal alpha function, Default=0 + * alpha_y: vertical alpha function, Default=0 + * Qpx: horizontal linear chromaticity, Default=0 + * Qpy: vertical linear chromaticity, Default=0 + * A1: horizontal amplitude detuning coefficient, Default=0 + * A2: cross term for amplitude detuning coefficient, Default=0 + * A3: vertical amplitude detuning coefficient, Default=0 + * emit_x: horizontal equilibrium emittance [m.rad], Default=0 + ignored if emit_x=0 + * emit_y: vertical equilibrium emittance [m.rad], Default=0 + ignored if emit_y=0 + * sigma_dp: equilibrium momentum spread, Default=0 + ignored if sigma_dp=0 + * tau_x: horizontal radiation damping time [turns], Default=0 + ignored if tau_x=0 + * tau_y: vertical radiation damping time [turns], Default=0 + ignored if tau_y=0 + * tau_z: longitudinal radiation damping time [turns], Default=0 + ignored if tau_z=0 + * U0: - energy loss [eV] (positive number), Default=0 + * TimeLag: Set the timelag of the cavities, Default=0. Can be scalar + or sequence of scalars (as with harmonic_number and Vrf). + + If the given emit_x,emit_y or sigma_dp is 0, then no equlibrium emittance + is applied in this plane. + If the given tau is 0, then no radiation damping is applied for this plane. + + + Returns: + ring (Lattice): Simple ring + """ + + # compute slip factor + gamma = energy / e_mass + eta = alpha - 1/gamma**2 + + harmonic_number = numpy.atleast_1d(harmonic_number) + Vrf = numpy.atleast_1d(Vrf) + try: + TimeLag = numpy.broadcast_to(TimeLag, Vrf.shape) + except ValueError: + raise AtError('TimeLag needs to be broadcastable to Vrf (same shape)') + + if (len(harmonic_number) != len(Vrf)): + raise AtError('harmonic_number input must match length of Vrf input') + + # compute rf frequency + frf = harmonic_number * clight / circumference + + all_cavities = [] + for icav in numpy.arange(len(harmonic_number)): + # generate rf cavity element + rfcav = RFCavity('RFC', 0, Vrf[icav], frf[icav], harmonic_number[icav], + energy, TimeLag=TimeLag[icav]) + all_cavities.append(rfcav) + + # Now we will use the optics parameters to compute the uncoupled M66 matrix + + s_dphi_x = numpy.sin(2*numpy.pi*Qx) + c_dphi_x = numpy.cos(2*numpy.pi*Qx) + s_dphi_y = numpy.sin(2*numpy.pi*Qy) + c_dphi_y = numpy.cos(2*numpy.pi*Qy) + + M00 = c_dphi_x + alpha_x * s_dphi_x + M01 = beta_x * s_dphi_x + M10 = -(1. + alpha_x**2) / beta_x * s_dphi_x + M11 = c_dphi_x - alpha_x * s_dphi_x + + M22 = c_dphi_y + alpha_y * s_dphi_y + M23 = beta_y * s_dphi_y + M32 = -(1. + alpha_y**2) / beta_y * s_dphi_y + M33 = c_dphi_y - alpha_y * s_dphi_y + + 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.], + [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) + lin_elem = M66('Linear', m66=Mat66, Length=circumference) + + # Generate the simple quantum diffusion element + quantdiff = SimpleQuantDiff('SQD', beta_x=beta_x, beta_y=beta_y, + emit_x=emit_x, emit_y=emit_y, + sigma_dp=sigma_dp, tau_x=tau_x, + tau_y=tau_y, tau_z=tau_z, U0=U0) + + # Generate the detuning element + nonlin_elem = Element('NonLinear', PassMethod='DeltaQPass', + Betax=beta_x, Betay=beta_y, + Alphax=alpha_x, Alphay=alpha_y, + Qpx=Qpx, Qpy=Qpy, + A1=A1, A2=A2, A3=A3) + + # Assemble all elements into the lattice object + ring = Lattice(all_cavities + [lin_elem, nonlin_elem, quantdiff], + energy=energy, periodicity=1) + + return ring diff --git a/pyat/test/test_physics.py b/pyat/test/test_physics.py index 27059f353..28d66c858 100644 --- a/pyat/test/test_physics.py +++ b/pyat/test/test_physics.py @@ -1,6 +1,7 @@ import at import numpy from numpy.testing import assert_allclose as assert_close +from numpy.testing import assert_equal import pytest from at import AtWarning, physics from at import lattice_track @@ -319,7 +320,14 @@ def test_quantdiff(hmba_lattice): 0.00000000e+00, 3.71123417e-14, 5.61789810e-17]], rtol=1e-5, atol=1e-20) - +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(ring[-1].PassMethod, 'SimpleQuantDiffPass') + ring.disable_6d() + assert_equal(ring[-1].PassMethod, 'IdentityPass') + assert_close(ring.get_tune(), [0.1,0.2], atol=1e-10) + @pytest.mark.parametrize('refpts', ([121], [0, 40, 121])) def test_ohmi_envelope(hmba_lattice, refpts): hmba_lattice = hmba_lattice.radiation_on(quadrupole_pass=None, copy=True)