Collective unit tests (#539)

* 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 * add collective tests * fix imports * add tracking test * mode=wake * mode=wake * check vals wake Co-authored-by: Lee Carver <[email protected]> Co-authored-by: Lee Carver <[email protected]>
atcollab · Dec 12, 2022 · 39318eb · 39318eb
1 parent abb815a
commit 39318eb
Showing 1 changed file with 136 additions and 0 deletions.
diff --git a/pyat/test/test_collective.py b/pyat/test/test_collective.py
@@ -0,0 +1,136 @@
+import at
+import numpy
+import pytest
+from numpy.testing import assert_allclose as assert_close
+from numpy.testing import assert_equal
+from at.collective import Wake, WakeElement, ResonatorElement
+from at.collective import WakeComponent, ResWallElement
+from at.collective import add_beamloading, remove_beamloading, BLMode
+
+
+_issorted = lambda a: numpy.all(a[:-1] <= a[1:])
+
+
+def test_fillpattern(hmba_lattice):
+    ring = hmba_lattice.radiation_on(copy=True)
+    nbunch = 16
+    current = 0.2
+    ring.set_fillpattern(nbunch)
+    ring.beam_current = current
+    assert ring.fillpattern.shape == (ring.harmonic_number, )
+    assert numpy.sum(ring.fillpattern) == 1
+    assert ring.bunch_list.shape == (nbunch, )
+    assert ring.nbunch == nbunch
+    assert ring.bunch_currents.shape == (nbunch, )
+    assert numpy.sum(ring.bunch_currents) == current
+
+
+def test_build_srange():
+    bspacing = 10
+    wlength = 20
+    bsize = 0.1
+    steps = 1.0e-3
+    stepl = 1.0e-2
+    srange = Wake.build_srange(-bsize, bsize, steps, stepl, bspacing, wlength)
+    assert len(numpy.unique(srange)) == len(srange)
+    assert _issorted(srange)
+    assert_close(max(srange), wlength + bsize - bspacing - stepl)
+    assert min(srange) == -bsize
+    assert_close(min(numpy.diff(srange)) ,steps)
+    assert_close(min(numpy.diff(srange[srange>bsize])) ,stepl)
+
+
+def test_wake_object():
+    srange = [0, 1, 1]
+    with pytest.raises(Exception):
+        Wake(srange)
+    srange = Wake.build_srange(-0.36, 0.36, 1.0e-5, 1.0e-2, 844, 844)
+    with pytest.raises(Exception):
+        Wake.resistive_wall(srange, Wake.WakeComponent.X, 1.0, 1.0e-3, 1.0)
+    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)
+    assert 1.0e-24 in long_res.srange
+    assert long_res.Z.shape == long_res.srange.shape
+
+
+def test_wake_element(hmba_lattice):
+    ring = hmba_lattice.radiation_on(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)
+    assert welem.TurnHistory.shape == (ring.nbunch*welem.Nslice*welem.Nturns, 4)
+    ring.append(welem)
+    ring.set_fillpattern(16)
+    assert welem.TurnHistory.shape == (ring.nbunch*welem.Nslice*welem.Nturns, 4)
+    ring.disable_6d(at.Collective)
+    assert welem.PassMethod == 'IdentityPass'
+
+
+def test_resonator_element(hmba_lattice):
+    ring = hmba_lattice.radiation_on(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)
+    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)
+    with pytest.raises(Exception):
+        add_beamloading(ring, 44e3, 400, cavpts=range(len(ring)))
+    add_beamloading(ring, 44e3, 400)
+    cavs = ring.get_elements(at.RFCavity)  
+    for cav in cavs:
+        assert cav.PassMethod == 'BeamLoadingCavityPass'
+        assert hasattr(cav, 'Vbeam') 
+        assert hasattr(cav, 'Vgen')
+        assert hasattr(cav, 'Vcav') 
+    ring.disable_6d(at.RFCavity)
+    for cav in cavs:
+        assert cav.PassMethod == 'IdentityPass'  
+    ring.enable_6d(at.RFCavity)
+    for cav in cavs:
+        assert cav.PassMethod == 'BeamLoadingCavityPass'          
+    remove_beamloading(ring)
+    cavs = ring.get_elements(at.RFCavity)  
+    for cav in cavs:
+        assert cav.PassMethod == 'RFCavityPass' 
+
+
+def test_track_beamloading(hmba_lattice):
+    ring = hmba_lattice.radiation_on(copy=True)
+    rin0 = numpy.zeros(6)
+    at.lattice_pass(ring, rin0, refpts=[])
+    add_beamloading(ring, 44e3, 400, mode=BLMode.WAKE)
+    rin1 = numpy.zeros(6)
+    at.lattice_pass(ring, rin1, refpts=[])    
+    assert_close(rin0, rin1, atol=1e-21)
+    ring.set_fillpattern(2)
+    ring.beam_current = 0.2
+    rin = numpy.zeros((6,1))
+    with pytest.raises(Exception):
+        at.lattice_pass(ring, rin, refpts=[])
+    rin = numpy.zeros((6,2))
+    at.lattice_pass(ring, rin, refpts=[])
+    assert_close(rin[:,0], numpy.array([-2.318948e-08, -1.599715e-09,
+                                        0.000000e+00,  0.000000e+00,
+                                        -1.313306e-05, -1.443748e-08]
+                                        ), atol=1e-10)
+