Regression bug from PR atcollab#774: beam_current lattice attribute n…

…ot propagated correctly (atcollab#778) * correct bug, improve help and fix pep8 * refactor eq and multiply by 4 second term * properly copy over beam_current * resotre old rotate_table * fix memmove * fix memmove * single loop onz * remove memmove * add test * ignore warning * change to wake * ignore warning * add vbeam buffers * Update example scripts to include disbale/enable 6d * use enable_6d intests --------- Co-authored-by: Lee Carver <[email protected]>
T-Nicholls · Jun 20, 2024 · 131f0b2 · 131f0b2
1 parent 201e57a
commit 131f0b2
Show file tree

Hide file tree

Showing 8 changed files with 95 additions and 35 deletions.
diff --git a/atintegrators/BeamLoadingCavityPass.c b/atintegrators/BeamLoadingCavityPass.c
@@ -43,7 +43,7 @@ struct elem
 
 void write_buffer(double *data, double *buffer, int datasize, int buffersize){
     if(buffersize>1){
-        memmove(buffer, buffer + datasize, datasize*buffersize*sizeof(double));
+        memmove(buffer, buffer + datasize, datasize*(buffersize-1)*sizeof(double));
     }
     memcpy(buffer + datasize*(buffersize-1), data, datasize*sizeof(double));
 }

diff --git a/atintegrators/atimplib.c b/atintegrators/atimplib.c
@@ -38,23 +38,35 @@ static double getTableWake(double *waketable,double *waketableT,double distance,
 };
 
 static void rotate_table_history(long nturns,long nslice,double *turnhistory,double circumference){
-    int i;
-    if(nturns > 1){
-        memmove(turnhistory, turnhistory + nslice, 4*nslice*nturns*sizeof(double));
-        double *z = turnhistory+nslice*nturns*2;
-        for(i=0; i<nslice*nturns; i++){
-            z[i] += -circumference;
+
+    double *xtmp,*xtmp0;
+    double *ytmp,*ytmp0;
+    double *ztmp,*ztmp0;
+    double *wtmp,*wtmp0;
+    double *t0;
+    int i, ii;    
+    for (i=0;i<nturns-1;i++){
+        xtmp0 = turnhistory + i*nslice;
+        xtmp = turnhistory + (i+1)*nslice;
+        ytmp0 = turnhistory + (i+nturns)*nslice;
+        ytmp = turnhistory + (i+nturns+1)*nslice;
+        ztmp0 = turnhistory + (i+2*nturns)*nslice;
+        ztmp = turnhistory + (i+2*nturns+1)*nslice;
+        wtmp0 = turnhistory + (i+3*nturns)*nslice;
+        wtmp = turnhistory + (i+3*nturns+1)*nslice;
+        for(ii=0;ii<nslice;ii++){
+            xtmp0[ii]=xtmp[ii];
+            ytmp0[ii]=ytmp[ii];
+            ztmp0[ii]=ztmp[ii]-circumference;
+            wtmp0[ii]=wtmp[ii];
         }
     }
-    double *x0 = turnhistory + (nturns-1)*nslice;
-    double *y0 = turnhistory + (2*nturns-1)*nslice;
-    double *z0 = turnhistory + (3*nturns-1)*nslice;
-    double *w0 = turnhistory + (4*nturns-1)*nslice;
-    for(i=0; i<nslice; i++){
-        x0[i] = 0.0;
-        y0[i] = 0.0;
-        z0[i] = 0.0;
-        w0[i] = 0.0;
+
+    for(ii=1;ii<5; ii++){ 
+        t0 = turnhistory + (ii*nturns-1)*nslice;
+        for(i=0; i<nslice; i++){
+            t0[i] = 0.0;
+        }
     }
 };
 

diff --git a/pyat/at/lattice/lattice_object.py b/pyat/at/lattice/lattice_object.py
@@ -93,7 +93,7 @@ class Lattice(list):
     _excluded_attributes = ('nbunch', )
     # Attributes propagated in copies:
     _std_attributes = ('name', '_energy', '_particle', 'periodicity',
-                       '_cell_harmnumber', '_radiation', 'beam_current',
+                       '_cell_harmnumber', '_radiation', '_beam_current',
                        '_fillpattern')
 
     # noinspection PyUnusedLocal

diff --git a/pyat/examples/CollectiveEffects/BeamLoading.py b/pyat/examples/CollectiveEffects/BeamLoading.py
@@ -52,7 +52,7 @@ def analytical_qs(ring, I):
 print(size, rank)
 
 ring = at.load_lattice('../../../machine_data/esrf.m')
-ring.radiation_on(cavity_pass='RFCavityPass')
+ring.enable_6d(cavity_pass='RFCavityPass')
 ring.set_rf_frequency()
 
 
@@ -74,7 +74,7 @@ def analytical_qs(ring, I):
 
 # Here we specify whether we want to use PHASOR or WAKE
 # beam loading models.
-mode = 'PHASOR'
+mode = 'WAKE'
 if mode == 'WAKE':
     blm = BLMode.WAKE
 else:
@@ -108,6 +108,18 @@ def analytical_qs(ring, I):
     z_all = numpy.zeros((Nturns, Nbunches))
     dp_all = numpy.zeros((Nturns, Nbunches))
 
+
+# Here it should be considered that there are 2 ways to run
+# simulations in pyat. Either with ring.track(part, nturns=Nturns)
+# or with for i in np.arange(Nturns): ring.track(part, nturns=1).
+# With the former, you should use the BeamMoments element to acquire
+# the means and stds of each bunch turn by turn when using MPI.
+# However, when you want to kick after a certain turn, you have 1
+# of 2 choices. Either you use a BeamMoments element, and split the 
+# tracking into 2 (before and after). You can then concatenate the
+# the results. Or you can use the code below to gather all particles
+# yourself with the MPICOMM after each turn.
+
 for i in numpy.arange(Nturns):
     # Apply a kick to ensure the coherent tune is large
     if i == kickTurn:

diff --git a/pyat/examples/CollectiveEffects/LCBI_run.py b/pyat/examples/CollectiveEffects/LCBI_run.py
@@ -48,11 +48,10 @@ def launch():
     ring.beam_current = current
     ring.set_rf_voltage(6e6)
 
-    ring.radiation_off(cavity_pass='RFCavityPass')
+    ring.disable_6d(cavity_pass='RFCavityPass')
     ring.set_fillpattern(Nbunches)
 
     ring.set_rf_frequency()
-    freqres = ring.rf_frequency
 
     qx, qy, qs = ring.get_tune()
     fs0 = qs * ring.revolution_frequency
@@ -77,7 +76,7 @@ def launch():
     fring.append(bmon)
 
     # Particle generation and tracking
-    sigm = at.sigma_matrix(ring.radiation_on(copy=True))
+    sigm = at.sigma_matrix(ring.enable_6d(copy=True))
     part = at.beam(Npart, sigm, orbit=ld0)
     part[4, :] = 0
     part[5, :] = 0

diff --git a/pyat/examples/CollectiveEffects/LongDistribution.py b/pyat/examples/CollectiveEffects/LongDistribution.py
@@ -13,7 +13,7 @@
 qfactor = 1
 Rs = 1e4
 current = 5e-4
-m = 50  # 30 is quite coarse, 70 or 80 is very fine. 50 is middle
+m = 60  # 30 is quite coarse, 70 or 80 is very fine. 50 is middle
 kmax = 8
 
 
@@ -60,7 +60,7 @@
 ring.beam_current = current
 ring.set_fillpattern(Nbunches)
 
-ring.radiation_on()
+ring.enable_6d()
 ring.set_cavity_phase()
 _, fring = at.fast_ring(ring)
 fring.append(welem)

diff --git a/pyat/examples/CollectiveEffects/TRW_run.py b/pyat/examples/CollectiveEffects/TRW_run.py
@@ -36,7 +36,7 @@ def launch():
 
     # Set up the ring
     ring = at.load_m('../../../machine_data/esrf.m')
-    ring.radiation_off(cavity_pass='RFCavityPass')
+    ring.disable_6d(cavity_pass='RFCavityPass')
 
     current = 200e-3
     nturns = 10000
@@ -85,7 +85,7 @@ def launch():
     fring.append(bmon)
 
     # Particle generation and tracking
-    sigm = at.sigma_matrix(ring.radiation_on(copy=True))
+    sigm = at.sigma_matrix(ring.enable_6d(copy=True))
     part = at.beam(Npart, sigm)
     part[0, :] = 1e-6
     part[1, :] = 0

diff --git a/pyat/test/test_collective.py b/pyat/test/test_collective.py
@@ -1,6 +1,7 @@
 import at
 import numpy
 import pytest
+import warnings
 from numpy.testing import assert_allclose as assert_close
 from at.collective import Wake, WakeElement, ResonatorElement
 from at.collective import WakeComponent, ResWallElement
@@ -13,7 +14,7 @@
 
 
 def test_fillpattern(hmba_lattice):
-    ring = hmba_lattice.radiation_on(copy=True)
+    ring = hmba_lattice.enable_6d(copy=True)
     nbunch = 16
     current = 0.2
     ring.set_fillpattern(nbunch)
@@ -55,7 +56,7 @@ def test_wake_object():
 
 
 def test_wake_element(hmba_lattice):
-    ring = hmba_lattice.radiation_on(copy=True)
+    ring = hmba_lattice.enable_6d(copy=True)
     srange = Wake.build_srange(0.0, 0.36, 1.0e-5, 1.0e-2, 844, 844) 
     long_res =  Wake.long_resonator(srange, 1.0e9, 1.0, 1.0e3, 1.0) 
     welem = WakeElement('WELEM', ring, long_res)
@@ -68,32 +69,30 @@ def test_wake_element(hmba_lattice):
 
 
 def test_resonator_element(hmba_lattice):
-    ring = hmba_lattice.radiation_on(copy=True)
+    ring = hmba_lattice.enable_6d(copy=True)
     srange = Wake.build_srange(0.0, 0.36, 1.0e-5, 1.0e-2, 844, 844)
     welem = ResonatorElement('WELEM', ring, srange, WakeComponent.Z, 1.0e9, 1.0, 1.0e3)  
     assert welem.ResFrequency == 1.0e9
     assert welem.Qfactor == 1
     assert welem.Rshunt == 1.0e3
-    print(welem.WakeZ[[0,1]])
     assert_close(welem.WakeZ[[0,1]], [3.14159265e+12, 6.28318531e+12], rtol = 1e-8)
     welem.ResFrequency += 100
     assert welem.ResFrequency == 1.0e9 + 100
 
 
 def test_resistive_wall_element(hmba_lattice):
-    ring = hmba_lattice.radiation_on(copy=True)
+    ring = hmba_lattice.enable_6d(copy=True)
     srange = Wake.build_srange(0.0, 0.36, 1.0e-5, 1.0e-2, 844, 844)
     welem = ResWallElement('WELEM', ring, srange, WakeComponent.DX, 1.0, 1.0e-2, 1.0e6)  
     assert welem.RWLength == 1.0
     assert welem.Conductivity == 1.0e6
-    print(welem.WakeDX[[0,1]])
     assert_close(welem.WakeDX[[0,1]], [ 0.0000000e+00, -1.0449877e+24], rtol = 1e-8)
     welem.Conductivity += 100
     assert welem.Conductivity == 1.0e6 + 100
 
 
 def test_beamloading(hmba_lattice):
-    ring = hmba_lattice.radiation_on(copy=True)
+    ring = hmba_lattice.enable_6d(copy=True)
     with pytest.raises(Exception):
         add_beamloading(ring, 44e3, 400, cavpts=range(len(ring)))
     add_beamloading(ring, 44e3, 400)
@@ -117,7 +116,7 @@ def test_beamloading(hmba_lattice):
 
 @pytest.mark.parametrize('func', (lattice_track, lattice_pass))
 def test_track_beamloading(hmba_lattice, func):
-    ring = hmba_lattice.radiation_on(copy=True)
+    ring = hmba_lattice.enable_6d(copy=True)
     rin0 = numpy.zeros(6)
     func(ring, rin0, refpts=[])
     add_beamloading(ring, 44e3, 400, blmode=BLMode.WAKE)
@@ -138,4 +137,42 @@ def test_track_beamloading(hmba_lattice, func):
                                         0.000000e+00,  0.000000e+00,
                                         -1.313306e-05, -1.443748e-08]
                                         ), atol=5e-10)
-
+
+
+def test_buffers(hmba_lattice):
+    ring = hmba_lattice.enable_6d(copy=True)
+    ring.periodicity = 1
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        ring.harmonic_number = 32
+    nturns = 11
+    nbunch = 4
+    nslice = 51
+    ns = nbunch*nslice
+    ls = ns*ring.circumference/ring.periodicity
+    add_beamloading(ring, 44e3, 400, Nturns=nturns, Nslice=nslice,
+                    buffersize=nturns, blmode=BLMode.WAKE)
+    ring.set_fillpattern(nbunch)
+    ring.beam_current = 0.2
+    rin = numpy.zeros((6, nbunch)) + 1.0e-6
+    bl_elem = ring.get_elements('*_BL')[0]
+    th = numpy.zeros((nturns, ) + bl_elem.TurnHistory.shape)
+    vbh = numpy.zeros((nturns, ) + bl_elem.Vbeam_buffer.shape)
+    vgh = numpy.zeros((nturns, ) + bl_elem.Vgen_buffer.shape)
+    vbbh = numpy.zeros((nturns, ) + bl_elem.Vbunch_buffer.shape)
+    for i in numpy.arange(nturns):
+        ring.track(rin, nturns=1, refpts=None, in_place=True)
+        th[i] = bl_elem.TurnHistory
+        vbh[i] = bl_elem.Vbeam_buffer
+        vgh[i] = bl_elem.Vgen_buffer
+        vbbh[i] = bl_elem.Vbunch_buffer
+    for i in numpy.arange(1, nturns):
+        dth = numpy.sum(th[i-1, (nturns-i)*ns:]
+                        - th[i, (nturns-i-1)*ns:(nturns-1)*ns]) - i*ls
+        dvbh = numpy.sum(vbh[i-1, :, (nturns-i):]
+                         - vbh[i, :, (nturns-i-1):(nturns-1)])
+        dvgh = numpy.sum(vgh[i-1, :, (nturns-i):]
+                         - vgh[i, :, (nturns-i-1):(nturns-1)])
+        dvbbh = numpy.sum(vbbh[i-1, :, (nturns-i):]
+                          - vbbh[i, :, (nturns-i-1):(nturns-1)])
+    assert_close([dth, dvbh, dvgh, dvbbh], numpy.zeros(4), atol=1e-9)