diff --git a/.gitignore b/.gitignore index 2ec607fe88cb09dace70c5df9b331a9a54f82aff..50648b4d3fa0c764efbd72e16ad4cae1a31a5458 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,4 @@ *__pycache__* *.ipynb_checkpoints* -test_*.py *.hdf5 +*.pyc diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..cee5c6b2fd798e5e1caec32a80409be676f32d94 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,47 @@ +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 diff --git a/tests/test_bunch.py b/tests/test_bunch.py new file mode 100644 index 0000000000000000000000000000000000000000..19aceaf8043d18b17c2a9ba6774dec7b18aedac7 --- /dev/null +++ b/tests/test_bunch.py @@ -0,0 +1,71 @@ +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 diff --git a/tests/test_imports.py b/tests/test_imports.py new file mode 100644 index 0000000000000000000000000000000000000000..11a9edfb665ae6c357eb2635d5e9602ce17bb7ac --- /dev/null +++ b/tests/test_imports.py @@ -0,0 +1,5 @@ +from mbtrack2 import * + + +def test_imports(): + assert True \ No newline at end of file diff --git a/tests/test_optics.py b/tests/test_optics.py new file mode 100644 index 0000000000000000000000000000000000000000..288915494955760da047420f76ca0accbe9cffbf --- /dev/null +++ b/tests/test_optics.py @@ -0,0 +1,15 @@ +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 diff --git a/tests/test_synchrotron.py b/tests/test_synchrotron.py new file mode 100644 index 0000000000000000000000000000000000000000..32572aba59a15e896e2b788c9eb534e07156894f --- /dev/null +++ b/tests/test_synchrotron.py @@ -0,0 +1,61 @@ +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