test_bunch.py

import numpy as np
import pytest
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