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 <[email protected]> Co-authored-by: Lee Carver <[email protected]>
atcollab · Aug 27, 2023 · 8b8fca8 · 8b8fca8
1 parent 1d8cd5b
commit 8b8fca8
Show file tree

Hide file tree

Showing 5 changed files with 403 additions and 14 deletions.
diff --git a/atintegrators/SimpleQuantDiffPass.c 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; c<num_particles; c++) { /*Loop over particles  */
+    r6 = r_in+c*6;
+    double randnorm[3];
+    for (i = 0; i < 3; i++) {
+        randnorm[i] = atrandn_r(rng, 0.0, 1.0);
+    }
+
+    if(!atIsNaN(r6[0])) {
+      if(tau_x!=0.0) {
+        r6[1] -= 2*r6[1]/tau_x;
+      }
+      if(sigma_xp!=0.0) {
+        r6[1] += 2*sigma_xp*sqrt(1/tau_x)*randnorm[0];
+      }
+      if(sigma_yp!=0.0) {
+        r6[3] += 2*sigma_yp*sqrt(1/tau_y)*randnorm[1];
+      }
+      if(tau_y!=0.0) {
+        r6[3] -= 2*r6[3]/tau_y;
+      }
+      if(sigma_dp!=0.0) {
+        r6[4] += 2*sigma_dp*sqrt(1/tau_z)*randnorm[2];
+      }
+      if(tau_z!=0.0) {
+        r6[4] -= 2*r6[4]/tau_z;
+      }
+
+      if(EnergyLossFactor>0.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
@@ -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
@@ -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')
     )
 }