Unit tests for Synchrotron, Optics and Bunch

.gitignore

 *__pycache__*
 *.ipynb_checkpoints*
-test_*.py
 *.hdf5
+*.pyc

tests/conftest.py

+import pytest
+import numpy as np
+from mbtrack2 import Optics, Synchrotron, Electron, Bunch
+
+@pytest.fixture
+def local_optics():
+    beta = np.array([1, 1])
+    alpha = np.array([0, 0])
+    dispersion = np.array([0, 0, 0, 0])
+    local_optics = Optics(local_beta=beta, local_alpha=alpha, 
+                      local_dispersion=dispersion)
+    return local_optics
+
+@pytest.fixture
+def demo_ring(local_optics):
+    
+    h = 20
+    L = 100
+    E0 = 1e9
+    particle = Electron()
+    ac = 1e-3
+    U0 = 250e3
+    tau = np.array([10e-3, 10e-3, 5e-3])
+    tune = np.array([18.2, 10.3])
+    emit = np.array([50e-9, 50e-9*0.01])
+    sigma_0 = 30e-12
+    sigma_delta = 1e-3
+    chro = [1.0,1.0]
+    
+    ring = Synchrotron(h, local_optics, particle, L=L, E0=E0, ac=ac, U0=U0, tau=tau,
+                       emit=emit, tune=tune, sigma_delta=sigma_delta, 
+                       sigma_0=sigma_0, chro=chro)
+    
+    return ring
+
+@pytest.fixture
+def mybunch(demo_ring):
+    mp_number = 10
+    mybunch = Bunch(demo_ring, mp_number=mp_number, track_alive=True)
+    return mybunch
+
+@pytest.fixture
+def large_bunch(demo_ring):
+    mp_number = 1e5
+    large_bunch = Bunch(demo_ring, mp_number=mp_number, track_alive=True)
+    large_bunch.init_gaussian()
+    return large_bunch
\ No newline at end of file

tests/test_bunch.py

+import pytest
+import numpy as np
+from scipy.constants import e
+from mbtrack2 import Bunch
+
+def test_bunch_values(demo_ring):
+    mp_number = 10
+    current = 20e-3
+    mybunch = Bunch(demo_ring, mp_number=mp_number, current=current, 
+                    track_alive=True)
+    
+    assert mybunch.mp_number == mp_number
+    assert pytest.approx(mybunch.current) == current
+    assert len(mybunch) == mp_number
+    np.testing.assert_allclose(mybunch.alive, np.ones((mp_number,), dtype=bool))
+    assert pytest.approx(mybunch.charge) == current * demo_ring.T0
+    assert pytest.approx(mybunch.charge_per_mp) == current * demo_ring.T0 / mp_number
+    assert pytest.approx(mybunch.particle_number) == current * demo_ring.T0 / e
+    assert mybunch.is_empty == False
+    
+def test_bunch_magic(mybunch):
+    for label in mybunch:
+        np.testing.assert_allclose(mybunch[label], np.zeros(len(mybunch)))
+        mybunch[label] = np.ones(len(mybunch))
+        np.testing.assert_allclose(mybunch[label], np.ones(len(mybunch)))
+    
+def test_bunch_losses(mybunch):
+    charge_init = mybunch.charge
+    mybunch.alive[0] = False
+    assert len(mybunch) == mybunch.mp_number - 1
+    assert pytest.approx(mybunch.charge) == charge_init * len(mybunch) / mybunch.mp_number
+    
+def test_bunch_init_gauss(large_bunch):
+    large_bunch.init_gaussian(mean=np.ones((6,)))
+    np.testing.assert_allclose(large_bunch.mean, np.ones((6,)), rtol=1e-2)
+    
+def test_bunch_save_load(mybunch, demo_ring, tmp_path):
+    mybunch["x"] += 1
+    mybunch.save(str(tmp_path / "test"))
+    
+    mybunch2 = Bunch(demo_ring, mp_number=1, current=1e-5)
+    mybunch2.load(str(tmp_path / "test.hdf5"))
+    
+    assert mybunch.mp_number == mybunch2.mp_number
+    assert pytest.approx(mybunch.charge) == mybunch2.charge
+    for label in mybunch:
+        np.testing.assert_allclose(mybunch[label], mybunch2[label])
+        
+def test_bunch_stats(demo_ring, large_bunch):
+    large_bunch.init_gaussian()
+    np.testing.assert_array_almost_equal(large_bunch.mean, np.zeros((6,)), decimal=5)
+    sig = np.concatenate((demo_ring.sigma(), [demo_ring.sigma_0, demo_ring.sigma_delta]))
+    np.testing.assert_allclose(large_bunch.std, sig, rtol=1e-2)
+    np.testing.assert_allclose(large_bunch.emit[:2], demo_ring.emit, rtol=1e-2)
+    np.testing.assert_allclose(large_bunch.cs_invariant[:2], demo_ring.emit*2, rtol=1e-2)
+    
+def test_bunch_binning(mybunch):
+    mybunch.init_gaussian()
+    (bins, sorted_index, profile, center) = mybunch.binning()
+    profile0 = np.zeros((len(bins)-1,))
+    for i, val in enumerate(sorted_index):
+        assert bins[val] <= mybunch["tau"][i] <= bins[val+1]
+        profile0[val] += 1
+    np.testing.assert_allclose(profile0, profile)
+    
+def test_bunch_plots(mybunch):
+    mybunch.init_gaussian()
+    mybunch.plot_phasespace()
+    mybunch.plot_profile()
+    assert True
+    
\ No newline at end of file

tests/test_imports.py

+from mbtrack2 import *
+
+
+def test_imports():
+    assert True
\ No newline at end of file

tests/test_optics.py

+import pytest
+import numpy as np
+from mbtrack2 import Optics
+
+def test_local_optics():
+    beta = np.array([1, 1])
+    alpha = np.array([0, 0])
+    dispersion = np.array([0, 0, 0, 0])
+    gamma = (1 + alpha**2) / beta
+    optics = Optics(local_beta=beta, local_alpha=alpha, 
+                      local_dispersion=dispersion)
+    np.testing.assert_allclose(optics.local_beta, beta)
+    np.testing.assert_allclose(optics.local_alpha, alpha)
+    np.testing.assert_allclose(optics.local_dispersion, dispersion)
+    np.testing.assert_allclose(optics.local_gamma, gamma)
\ No newline at end of file

tests/test_synchrotron.py

+import pytest
+import numpy as np
+from scipy.constants import c, e
+from mbtrack2 import Synchrotron, Electron
+
+def test_synchrotron_values(local_optics):
+    h = 20
+    L = 100
+    E0 = 1e9
+    particle = Electron()
+    ac = 1e-3
+    U0 = 250e3
+    tau = np.array([10e-3, 10e-3, 5e-3])
+    tune = np.array([18.2, 10.3])
+    emit = np.array([50e-9, 50e-9*0.01])
+    sigma_0 = 30e-12
+    sigma_delta = 1e-3
+    chro = [1.0,1.0]
+    
+    ring = Synchrotron(h, local_optics, particle, L=L, E0=E0, ac=ac, U0=U0, tau=tau,
+                       emit=emit, tune=tune, sigma_delta=sigma_delta, 
+                       sigma_0=sigma_0, chro=chro)
+    assert pytest.approx(ring.h) == h
+    assert pytest.approx(ring.L) == L
+    assert pytest.approx(ring.E0) == E0
+    assert pytest.approx(ring.U0) == U0
+    assert pytest.approx(ring.ac) == ac
+    np.testing.assert_allclose(ring.tau, tau)
+    np.testing.assert_allclose(ring.tune, tune)
+    np.testing.assert_allclose(ring.emit, emit)
+    assert pytest.approx(ring.sigma_0) == sigma_0
+    assert pytest.approx(ring.sigma_delta) == sigma_delta
+    np.testing.assert_allclose(ring.chro, chro)
+    assert pytest.approx(ring.T0) == L/c
+    assert pytest.approx(ring.T1) == L/c/h
+    assert pytest.approx(ring.f0) == c/L
+    assert pytest.approx(ring.f1) == 1/(L/c/h)
+    assert pytest.approx(ring.omega0) == 2 * np.pi * c/L
+    assert pytest.approx(ring.omega1) == 2 * np.pi * 1/(L/c/h)
+    assert pytest.approx(ring.k1) == 2 * np.pi * 1/(L/c/h) / c
+    assert pytest.approx(ring.gamma) == E0 / (particle.mass * c**2 / e)
+    assert pytest.approx(ring.beta) == np.sqrt(1 - (E0 / (particle.mass * c**2 / e))**-2)
+    
+def test_synchrotron_mcf(demo_ring):
+    demo_ring.mcf_order = [5e-4, 1e-4, 1e-3]
+    assert pytest.approx(demo_ring.mcf(0.5)) == 5e-4*(0.5**2) + 1e-4*0.5 + 1e-3
+    assert pytest.approx(demo_ring.eta(0.5)) == demo_ring.mcf(0.5) - 1 / (demo_ring.gamma**2)
+    
+def test_synchrotron_tune(demo_ring):
+    tuneS = demo_ring.synchrotron_tune(1e6)
+    assert pytest.approx(tuneS, rel=1e-4) == 0.0017553
+    
+def test_synchrotron_long_twiss(demo_ring):
+    tuneS, long_alpha, long_beta, long_gamma = demo_ring.get_longitudinal_twiss(1e6, add=False)
+    assert pytest.approx(tuneS, rel=1e-4) == demo_ring.synchrotron_tune(1e6)
+    assert pytest.approx(long_alpha, rel=1e-4) == -0.0055146
+    assert pytest.approx(long_beta, rel=1e-4) == 3.0236e-08
+    assert pytest.approx(long_gamma, rel=1e-4) == 3.30736e7
+    
+def test_synchrotron_sigma(demo_ring):
+    np.testing.assert_allclose(demo_ring.sigma(), np.array([2.23606798e-04, 2.23606798e-04, 2.23606798e-05, 2.23606798e-05]))
\ No newline at end of file