Update examples (+ bug fix) (#527)

* Remove some unecessary prints * Fix bug in wake object resonator definition * Update file for modern usage * Added beam monitor and fixed Robinson file * Add srange value at 1e-24 to avoid interpolation issue in longresonator * Add beam monitor, update syntax * Add beam monitor * Added beamonitor * Adding the 1e-24 made this haissinski not work due to it taking the wrong ds, this is fixed * Pep8 * Pep8 * Pep8 * fix mexfunctions for beam loading * Add srange modifications to wake_functions * pep8 * Changed error * Moved srange modification to wake_object, with case for resonator or res wall * Modify memory for bspos and bcurr * Moved unique check to init of wake object and removed from resonator and reswall factory functions * Added assertion error for unique values of srange * Updated reswall function to include mask to allow summing of longres and trw into one element. Fixed hole in logic * correct condition for RW wfunc * wrong mode check for windows * Improved help for Haissinski * Add imports Co-authored-by: Lee Carver <[email protected]> Co-authored-by: Lee Carver <[email protected]> Co-authored-by: Simon White <[email protected]>
atcollab · Dec 9, 2022 · abb815a · abb815a
1 parent 6c2494b
commit abb815a
Show file tree

Hide file tree

Showing 12 changed files with 179 additions and 157 deletions.
diff --git a/atintegrators/BeamLoadingCavityPass.c b/atintegrators/BeamLoadingCavityPass.c
@@ -186,7 +186,7 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem,
         atWarning("Number of particles not a multiple of the number of bunches: uneven bunch load.");
     }
     #ifdef _MSC_VER
-    if(Elem->mode==1){
+    if(Elem->mode==2){
         atError("Beam loading Phasor mode not implemented in Windows.");
     }
     #endif
@@ -204,6 +204,12 @@ 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]);
+      struct elem El, *Elem=&El;
+
       long nslice,nturns,mode;
       double wakefact;
       double normfact, phasegain, voltgain;
@@ -266,13 +272,10 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
       /* ALLOCATE memory for the output array of the same size as the input  */
       plhs[0] = mxDuplicateArray(prhs[1]);
       r_in = mxGetDoubles(plhs[0]);
-      double *bspos = malloc(sizeof(double));
-      double *bcurr = malloc(sizeof(double));
-      bspos[0] = 0.0;
-      bcurr[0] = 0.0;
-      BeamLoadingCavityPass(r_in,num_particles,1,bspos,bcur,1,0,Elem);
-      free(bspos);
-      free(bcur);
+
+      double bspos = 0.0;
+      double bcurr = 0.0;
+      BeamLoadingCavityPass(r_in,num_particles,1,&bspos,&bcurr,1,0,Elem);
   }
   else if (nrhs == 0)
   {   /* return list of required fields */

diff --git a/atintegrators/BeamMomentsPass.c b/atintegrators/BeamMomentsPass.c
@@ -95,6 +95,12 @@ MODULE_DEF(BeamMomentsPass)        /* Dummy module initialisation */
 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]);
+        struct elem El, *Elem=&El;
+
         double *means;
         double *stds; 
         means=atGetDoubleArray(ElemData,"_means"); check_error();

diff --git a/pyat/at/collective/haissinski.py b/pyat/at/collective/haissinski.py
@@ -3,13 +3,32 @@
 from at.constants import clight, qe
 from scipy.interpolate import interp1d
 import time
+from at.collective import Wake
+from at.lattice import Lattice
 
 
 class Haissinski(object):
-    '''
+    r"""
     Class to find the longitudinal distribution in the presence
     of a short range wakefield.
 
+    Parameters: 
+    
+        wake_object:    class initialized with Wake that contains a 
+          longitudinal wake Z and equivalent srange.
+        ring:           a lattice object that is used to extract 
+         machine parameters
+
+    Keyword Args:
+        m (int):        the number of points in the full
+          distribution
+        kmax:           the min and max of the range of the distribution.
+          See equation 27. In units of [:math:`\sigma_z0`]
+        current:        bunch current.
+        numIters (int): the number of iterations
+        eps:            the convergence criteria.
+    
+    
     This class is a direct implementation of the following paper:
     "Numerical solution of the Haïssinski equation for the equilibrium
     state of  a stored electron beam", R. Warnock, K.Bane, Phys. Rev. Acc.
@@ -19,27 +38,17 @@ class Haissinski(object):
     The equation number of key formula are written next to the relevant
     function.
 
-    As input:
-    wake_object is a object that contains an srange and Z array.
-    ring is a ring instance which is needed for machine parameters
-    (sigma_l, sigma_e, etc)
-
-    m is the number of points in the full distribution that you want
-    kmax is the min and max of the range of the distribution.
-    See equation 27. In units of sigma_z0
-    current is the bunch current.
-    numIters is the number of iterations
-    eps is the convergence criteria.
-
-    the functions solve or solve_steps can be used after initialisation
+    The functions solve or solve_steps can be used after initialisation
+    An example usage can be found in:
+      at/pyat/examples/Collective/LongDistribution.py
 
     Future developments of this class:
         Adding LR wake or harmonic cavity as done at SOLEIL. Needs
         to be added WITH this class which is just for short range wake.
-    '''
+    """
 
-    def __init__(self, wake_object, ring, m=12, kmax=1,
-                 current=1e-4, numIters=10, eps=1e-10):
+    def __init__(self, wake_object: Wake, ring: Lattice, m: int = 12, kmax: float = 1,
+                 current: float = 1e-4, numIters: int = 10, eps: float = 1e-10):
 
         self.circumference = ring.circumference
         self.energy = ring.energy
@@ -61,7 +70,7 @@ def __init__(self, wake_object, ring, m=12, kmax=1,
 
         #  negative s to be consistent with paper and negative Wz
         s = wake_object.srange/self.sigma_l
-        self.ds = numpy.diff(s)[0]
+        self.ds = numpy.diff(s)[-1]
         self.s = -s
         self.wtot_fun = interp1d(self.s, -wake_object.Z,
                                  bounds_error=False, fill_value = 0)
@@ -103,8 +112,6 @@ def precompute_S(self):
         Equation 16
         '''
         print('Computing integrated wake potential')
-        print(self.q_array)
-        print(self.ds)
         sr = numpy.arange(2*numpy.amin(self.q_array),
                           numpy.abs(2*numpy.amin(self.q_array))
                           + self.ds, self.ds)

diff --git a/pyat/at/collective/wake_functions.py b/pyat/at/collective/wake_functions.py
@@ -10,12 +10,12 @@ def convolve_wakefun(srange, w, sigs):
     """Convolution of a wake function with a pulse of rms
     length sigs, this is use to generate a wake potential
     that can be added to the output of EM code like GDFIDL"""
-    sigt = sigs/clight
+    sigt = sigs / clight
     min_step = numpy.diff(srange)[0]
     t_out = numpy.arange(srange[0], srange[-1], min_step)
     sdiff = t_out[-1]-t_out[0]
     npoints = len(t_out)
-    nt = npoints+npoints-1
+    nt = npoints + npoints-1
     func = interp1d(srange, w, bounds_error=False, fill_value=0)
     wout = func(t_out)
     wout = numpy.append(wout, numpy.zeros(nt-len(wout)))
@@ -36,6 +36,7 @@ def long_resonator_wf(srange, frequency, qfactor, rshunt, beta):
     with the given parameters according to Alex Chao's resonator
     model (Eq. 2.82) and definitions of the resonator in HEADTAIL.
     """
+
     omega = 2 * numpy.pi * frequency
     alpha = omega / (2 * qfactor)
     omegabar = numpy.sqrt(numpy.abs(omega**2 - alpha**2))
@@ -87,17 +88,12 @@ def transverse_reswall_wf(srange, yokoya_factor, length, rvac, conduct, beta):
     HEADTAIL
     """
 
-    if numpy.amin(srange) <= 0:
-        raise ValueError("""
-                        Provided srange has either negative values or 0s
-                        This is not allowed for the transverse resistive wall
-                        wake function. Please correct.
-                        """)
-
     z0 = 119.9169832 * numpy.pi
     dt = -srange/(beta * clight)
     wake = (yokoya_factor * (numpy.sign(dt) - 1) / 2. *
             beta * length / numpy.pi / rvac**3 *
             numpy.sqrt(-z0 * clight / conduct / numpy.pi / dt))
-
+
+    wake[srange <= 0] = 0.0
+
     return wake
diff --git a/pyat/at/collective/wake_object.py b/pyat/at/collective/wake_object.py
@@ -62,6 +62,8 @@ class Wake(object):
     Components may be retrieved using :code:`wake.DX` for example
     """
     def __init__(self, srange):
+        assert len(srange) == len(numpy.unique(srange)), \
+            "srange must contain unique values"            
         self._srange = srange
         self.components = {WakeComponent.DX: None,
                            WakeComponent.DY: None,
@@ -141,10 +143,20 @@ def _readwakefile(self, filename, scol=0, wcol=1, sfact=1, wfact=1,
 
     def _resonator(self, wcomp, frequency, qfactor, rshunt, beta,
                    yokoya_factor=1):
-        if numpy.any(self._srange < 0):
-            warnings.warn(AtWarning('You are adding a wake function with '
-                                    'negative srange. This may cause issues '
-                                    'interpolating around zero.\n'))
+
+        if numpy.amin(self._srange) < 0:
+            raise ValueError("""
+                            Provided srange has negative values.
+                            This is not allowed for the resonator
+                            wake function. Please correct.
+                            """)
+
+        # It is needed to have a point at 0 and 1e-24 to sample properly
+        # the discontinuity in the longitudinal plane
+
+        self._srange = numpy.unique(numpy.concatenate(([0.0, 1e-24],
+                                                        self._srange)))
+
         if wcomp is WakeComponent.Z:
             return long_resonator_wf(self._srange, frequency,
                                      qfactor, rshunt, beta)
@@ -156,10 +168,20 @@ def _resonator(self, wcomp, frequency, qfactor, rshunt, beta,
             raise AtError('Invalid WakeComponent: {}'.format(wcomp))
 
     def _reswall(self, wcomp, length, rvac, conduct, beta, yokoya_factor=1):
-        if numpy.any(self._srange < 0):
-            warnings.warn(AtWarning('You are adding a wake function with '
-                                    'negative srange. This may cause issues '
-                                    'interpolating around zero.\n'))
+
+        if numpy.amin(self._srange) < 0:
+            raise ValueError("""
+                            Provided srange has negative values.
+                            This is not allowed for the resistive wall
+                            wake function. Please correct.
+                            """)
+
+        # It is needed to have a point at 0 and 1e-24 to sample properly
+        # the discontinuity in the longitudinal plane (if combining with
+        # longres)
+        self._srange = numpy.unique(numpy.concatenate(([0.0, 1e-24],
+                                                       self._srange)))
+
         if wcomp is WakeComponent.Z:
             raise AtError('Resitive wall not available '
                           'for WakeComponent: {}'.format(wcomp))
@@ -189,7 +211,7 @@ def resonator(srange, wakecomp,
         Several resonators can be built in one step by setting ``nelems``> 1
         and giving array parameters
         """
-        wake = Wake(numpy.unique(srange.clip(0)))
+        wake = Wake(srange)
         try:
             wakecomp = numpy.broadcast_to(wakecomp, (nelems, ))
             frequency = numpy.broadcast_to(frequency, (nelems, ))
@@ -231,8 +253,8 @@ def long_resonator(srange, frequency, qfactor, rshunt, beta, nelems=1):
             >>> Q = 1
             >>> Rs = 1e4
             >>> wake = Wake.long_resonator(srange, freq, Q, Rs, ring.beta)
-        """
-        return Wake.resonator(srange, WakeComponent.Z, rshunt, qfactor,
+        """ 
+        return Wake.resonator(srange, WakeComponent.Z, frequency, qfactor,
                               rshunt, beta, nelems=nelems)
 
     @staticmethod
@@ -251,7 +273,7 @@ def resistive_wall(srange, wakecomp: WakeComponent,
             yokoya_factor:  Yokoya factor
             nelems:
         """
-        wake = Wake(numpy.unique(srange.clip(0)))
+        wake = Wake(srange)
         try:
             wakecomp = numpy.broadcast_to(wakecomp, (nelems, ))
             length = numpy.broadcast_to(length, (nelems, ))

diff --git a/pyat/examples/CollectiveEffects/BeamLoading.py b/pyat/examples/CollectiveEffects/BeamLoading.py
@@ -9,8 +9,6 @@
 from at.physics import harmonic_analysis
 
 
-
-
 def analytical_qs(ring, I):
 
     E0 = ring.energy * qe
@@ -88,9 +86,9 @@ def analytical_qs(ring, I):
 # Now we give the fring and convert it
 # into a beam loaded cavity.
 add_beamloading(fring, qfactor, rshunt, Nslice=1,
-                          Nturns=50, mode=blm,
-                          VoltGain=0.1, PhaseGain=0.1)
-                          
+                Nturns=50, mode=blm,
+                VoltGain=0.1, PhaseGain=0.1)
+
 bl_elem = fring[at.get_refpts(fring, BeamLoadingElement)[0]]
 
 # Specify some simulation parameters
@@ -135,8 +133,8 @@ def analytical_qs(ring, I):
     qscoh = numpy.zeros(Nbunches)
     for ib in numpy.arange(Nbunches):
         qs = harmonic_analysis.get_tunes_harmonic(dp_all[kickTurn:, ib],
-                                                      num_harmonics=20,
-                                                      fmin=1e-5, fmax=0.1)
+                                                  num_harmonics=20,
+                                                  fmin=1e-5, fmax=0.1)
         qscoh[ib] = qs
 
     qs_mn, qs_std = numpy.array([numpy.mean(qscoh), numpy.std(qscoh)])

diff --git a/pyat/examples/CollectiveEffects/LCBI_analyse.py b/pyat/examples/CollectiveEffects/LCBI_analyse.py
@@ -115,30 +115,26 @@ def fitTimeVectors(tv):
 gr_peak = gr_theory[expected_mode]
 
 
-# Find the last turn in the data where there were no lost particles
-try:
-    tend = np.where(not resDict['lostturn'])[0][0]
-except:
-    tend = len(resDict['dp'])
-
 data_dp = resDict['dp'].T
 data_z = resDict['z'].T
 
 # Plot all data dp, just to see how it looks
 bbs = np.arange(0, Nbunches)
 for bb in bbs:
-    plt.plot(data_dp[bb, :tend])
+    plt.plot(data_dp[bb, :])
+plt.xlabel('Turn')
+plt.ylabel('dp/p')
 plt.show()
 
 # Here we use dp and z and combine with betaz into one complex number
 # this allows us to distinguish from + and - modes
 inv = np.array([[1/np.sqrt(betaz), 0],
                 [0, np.sqrt(betaz)]])
 
-data = np.zeros(data_dp[:, :tend].shape) + 1j * 0
+data = np.zeros(data_dp[:, :].shape) + 1j * 0
 for i in np.arange(data_dp.shape[0]):
-    dp = data_dp[i, :tend]
-    z = data_z[i, :tend]
+    dp = data_dp[i, :]
+    z = data_z[i, :]
     datpos, datmom = np.matmul(inv, [z, dp])
     data[i, :] = datpos + 1j * datmom
 

diff --git a/pyat/examples/CollectiveEffects/LCBI_run.py b/pyat/examples/CollectiveEffects/LCBI_run.py
@@ -70,9 +70,6 @@ def launch():
     srange = Wake.build_srange(0., 0.41, 1.0e-3, 1.0e-1,
                                bunchSpacing, ring.circumference*wturns)
 
-    # This is needed to fix an interpolation issue with LongResElem
-    srange = np.sort(np.concatenate((srange, [1e-24])))
-
     fres = ring.get_rf_frequency() - ring.revolution_frequency - fs0
     qfactor = 1e4
     rshunt = 2e6
@@ -81,33 +78,21 @@ def launch():
                                  rshunt, Nturns=wturns, Nslice=1)
     fring.append(welem)
 
+    bmon = at.BeamMoments('mon')
+    fring.append(bmon)
+
     # Particle generation and tracking
     sigm = at.sigma_matrix(ring.radiation_on(copy=True))
     part = at.beam(Npart, sigm, orbit=ld0)
     part[4, :] = 0
     part[5, :] = 0
 
-    if rank == 0:
-        z_all = np.zeros((nturns, Nbunches))
-        dp_all = np.zeros((nturns, Nbunches))
-        turnMsk = np.zeros(nturns, dtype=bool)
-
-    for i in np.arange(nturns):
-        _ = at.lattice_pass(fring, part)
-        allPartsg = comm.gather(part)
-        if rank == 0:
-            allPartsg = np.concatenate(allPartsg, axis=1)
-            stds, means = get_bunches_std_mean(allPartsg, Nbunches)
-            dp_all[i] = np.array([x[4] for x in means])
-            z_all[i] = np.array([x[5] for x in means])
-
-            sumNan = np.sum(np.isnan(allPartsg))
-            if sumNan > 0:
-                turnMsk[i] = True
-            print(i, sumNan, np.mean(np.abs(dp_all[i])))
+    _ = at.lattice_pass(fring, part, nturns=nturns)
 
+    dp_all = bmon.means[4, :, :]
+    z_all = bmon.means[5, :, :]
     if rank == 0:
-        outDict = pkl.dump({'dp': dp_all, 'z': z_all, 'lostturn': ~turnMsk,
+        outDict = pkl.dump({'dp': dp_all.T, 'z': z_all.T,
                             'ring': ring, 'fres': fres, 'qfactor': qfactor,
                             'rshunt': rshunt, 'M': Nbunches,
                             'current': current},