diff --git a/tests/physics/test_ibs_phys.py b/tests/physics/test_ibs_phys.py
index 026369073ae1e3ae080108437e64ca8e01735cd0..d2d24cf79e014194bbadaf2989ab79f2fbf3c230 100644
--- a/tests/physics/test_ibs_phys.py
+++ b/tests/physics/test_ibs_phys.py
@@ -1,5 +1,6 @@
import numpy as np
import pytest
+
pytestmark = pytest.mark.phys
from mbtrack2 import Electron, Synchrotron
from mbtrack2.tracking import (
@@ -12,6 +13,7 @@ from mbtrack2.tracking import (
)
from mbtrack2.utilities import Optics
+
@pytest.fixture
def model_ring():
h = 416 # Harmonic number of the accelerator.
diff --git a/tests/physics/test_lrrw_phys.py b/tests/physics/test_lrrw_phys.py
index d4c92898a9b8efc189e9d5fd0405dd7d703ca1ca..e2198720d0d3dffedc0d4e93edbd9600b9746218 100644
--- a/tests/physics/test_lrrw_phys.py
+++ b/tests/physics/test_lrrw_phys.py
@@ -1,14 +1,24 @@
import numpy as np
-from mbtrack2 import BeamMonitor, rwmbi_growth_rate
-from mbtrack2 import Beam, LongRangeResistiveWall, RFCavity, LongitudinalMap, TransverseMap
-from utility_test_growth_rate import get_growth_beam
import pytest
+from utility_test_growth_rate import get_growth_beam
+
+from mbtrack2 import (
+ Beam,
+ BeamMonitor,
+ LongitudinalMap,
+ LongRangeResistiveWall,
+ RFCavity,
+ TransverseMap,
+ rwmbi_growth_rate,
+)
+
pytestmark = pytest.mark.phys
+
class TestLongRangeResistiveWallPhys:
def test_llrw_growth_rate_chroma0(self, demo_ring, tmp_path):
-
+
ring = demo_ring
mp_num = int(1)
n_turns = int(5e3)
@@ -18,49 +28,60 @@ class TestLongRangeResistiveWallPhys:
Itot = 1.0
Vc = 1.0e6
name = str(tmp_path / 'save')
- ring.chro = [0,0]
-
+ ring.chro = [0, 0]
+
growth_rate = rwmbi_growth_rate(ring, Itot, radius, rho, "y")
-
- RF = RFCavity(ring, 1, Vc, np.arccos(ring.U0/Vc))
+
+ RF = RFCavity(ring, 1, Vc, np.arccos(ring.U0 / Vc))
long = LongitudinalMap(ring)
trans = TransverseMap(ring)
-
+
bb = Beam(ring)
- filling_pattern = np.ones((ring.h,))*Itot/ring.h
-
- bb.init_beam(filling_pattern, mp_per_bunch=mp_num, track_alive=False, mpi=False)
-
- rw = LongRangeResistiveWall(ring, beam=bb, length=ring.L, rho=rho,
- radius=radius, types=["Wydip"], nt=nt)
-
+ filling_pattern = np.ones((ring.h, )) * Itot / ring.h
+
+ bb.init_beam(filling_pattern,
+ mp_per_bunch=mp_num,
+ track_alive=False,
+ mpi=False)
+
+ rw = LongRangeResistiveWall(ring,
+ beam=bb,
+ length=ring.L,
+ rho=rho,
+ radius=radius,
+ types=["Wydip"],
+ nt=nt)
+
save_every = 10
- bbmon = BeamMonitor(h=ring.h, save_every=save_every, buffer_size=10,
- total_size=int(n_turns/save_every),
- file_name=name, mpi_mode=False)
-
+ bbmon = BeamMonitor(h=ring.h,
+ save_every=save_every,
+ buffer_size=10,
+ total_size=int(n_turns / save_every),
+ file_name=name,
+ mpi_mode=False)
+
for i in range(n_turns):
if (i % 1000 == 0):
print(i)
-
+
RF.track(bb)
long.track(bb)
trans.track(bb)
-
+
rw.track(bb)
bbmon.track(bb)
-
+
bbmon.close()
-
+
path = name + '.hdf5'
gr, gr_std = get_growth_beam(ring, path, 1e3, 5e3)
- rel_error = np.abs(growth_rate-gr)*100/growth_rate
- rel_error_std = np.abs(gr_std)*100/growth_rate
+ rel_error = np.abs(growth_rate - gr) * 100 / growth_rate
+ rel_error_std = np.abs(gr_std) * 100 / growth_rate
print(f"Relative error: {rel_error} % (std: {rel_error_std} %)")
assert rel_error < 2
-
- # def test_llrw_growth_rate_chroma0_average_beta(self,
- # generate_ring_with_at_lattice,
+
+ # def test_llrw_growth_rate_chroma0_average_beta(self,
+ # generate_ring_with_at_lattice,
# tmp_path):
# tracking_loc = np.random.rand()*350
# ring = generate_ring_with_at_lattice(tracking_loc=tracking_loc)
@@ -73,42 +94,42 @@ class TestLongRangeResistiveWallPhys:
# Vc = 2.8e6
# name = str(tmp_path / 'save')
# ring.chro = [0,0]
-
+
# growth_rate = rwmbi_growth_rate(ring, Itot, radius, rho, "y")
-
+
# RF = RFCavity(ring, 1, Vc, np.arccos(ring.U0/Vc))
# long = LongitudinalMap(ring)
# trans = TransverseMap(ring)
-
+
# bb = Beam(ring)
# filling_pattern = np.ones((ring.h,))*Itot/ring.h
-
+
# bb.init_beam(filling_pattern, mp_per_bunch=mp_num, track_alive=False, mpi=False)
-
- # rw = LongRangeResistiveWall(ring, beam=bb, length=ring.L, rho=rho,
+
+ # rw = LongRangeResistiveWall(ring, beam=bb, length=ring.L, rho=rho,
# radius=radius, types=["Wydip"], nt=nt)
-
+
# save_every = 10
- # bbmon = BeamMonitor(h=ring.h, save_every=save_every, buffer_size=10,
- # total_size=int(n_turns/save_every),
+ # bbmon = BeamMonitor(h=ring.h, save_every=save_every, buffer_size=10,
+ # total_size=int(n_turns/save_every),
# file_name=name, mpi_mode=False)
-
+
# for i in range(n_turns):
# if (i % 100 == 0):
# print(i)
-
+
# RF.track(bb)
# long.track(bb)
# trans.track(bb)
-
+
# rw.track(bb)
# bbmon.track(bb)
-
+
# bbmon.close()
-
+
# path = name + '.hdf5'
# gr, gr_std = get_growth_beam(ring, path, 500, 1e3)
# rel_error = np.abs(growth_rate-gr)*100/growth_rate
# rel_error_std = np.abs(gr_std)*100/growth_rate
# print(f"Relative error: {rel_error} % (std: {rel_error_std} %)")
- # assert rel_error < 2
\ No newline at end of file
+ # assert rel_error < 2
diff --git a/tests/test_imports.py b/tests/test_imports.py
index b02158fcc70017ae00b4684412b5e6660369a5f1..38956b58b1e4cc9428515c118dce084c21002cc8 100644
--- a/tests/test_imports.py
+++ b/tests/test_imports.py
@@ -1,6 +1,8 @@
import pytest
+
pytestmark = pytest.mark.unit
from mbtrack2 import *
+
def test_imports():
assert True
\ No newline at end of file
diff --git a/tests/unit/impedance/test_impedance_model.py b/tests/unit/impedance/test_impedance_model.py
index 8c593f2982b34551a22de73fd57f1a1b05e588d9..83d6a37172f439ee07145a5c2710d480d4855611 100644
--- a/tests/unit/impedance/test_impedance_model.py
+++ b/tests/unit/impedance/test_impedance_model.py
@@ -1,23 +1,21 @@
-
-from mbtrack2 import ImpedanceModel, ComplexData, WakeField
-import pandas as pd
import numpy as np
+import pandas as pd
-component_list = ComplexData.name_and_coefficients_table().columns.to_list()
+from mbtrack2 import ComplexData, ImpedanceModel, WakeField
+component_list = ComplexData.name_and_coefficients_table().columns.to_list()
import pytest
+
pytestmark = pytest.mark.unit
+
class TestImpedanceModel:
# Adding WakeField objects with unique names should succeed without errors
- @pytest.mark.parametrize('pos1,pos2', [([0],[1]), (0,1), (0.5,1.5)])
- def test_add_wakefield(self,
- ring_with_at_lattice,
- generate_wakefield,
- pos1,
- pos2):
+ @pytest.mark.parametrize('pos1,pos2', [([0], [1]), (0, 1), (0.5, 1.5)])
+ def test_add_wakefield(self, ring_with_at_lattice, generate_wakefield,
+ pos1, pos2):
model = ImpedanceModel(ring_with_at_lattice)
wakefield1 = generate_wakefield(name="wake1")
wakefield2 = generate_wakefield(name="wake2")
@@ -31,72 +29,100 @@ class TestImpedanceModel:
assert pos2 in model.positions
# Adding global WakeField objects should correctly append them to the globals list
- def test_add_global_wakefield(self, ring_with_at_lattice, generate_wakefield):
+ def test_add_global_wakefield(self, ring_with_at_lattice,
+ generate_wakefield):
model = ImpedanceModel(ring_with_at_lattice)
- global_wakefield = generate_wakefield(name = "global_wake")
+ global_wakefield = generate_wakefield(name="global_wake")
model.add_global(global_wakefield)
assert "global_wake" in model.globals_names
assert global_wakefield in model.globals
# Test weighted global sum with average_beta
- def test_global_wakefield_weighted(self, ring_with_at_lattice, generate_wakefield):
- ring_with_at_lattice.optics.local_beta = np.array([0.5,1])
- model = ImpedanceModel(ring_with_at_lattice, average_beta=np.array([4, 2]))
-
- global_wakefield = generate_wakefield(name = "global_wake",
- wake_function_params=[{'component_type': "xdip",
- 'function': np.array([1, 2])},
- {'component_type': "ydip",
- 'function': np.array([1, 2])}],
- impedance_params=[{'component_type': "xdip",
- 'function': np.array([1, 2])},
- {'component_type': "ydip",
- 'function': np.array([1, 2])}])
-
+ def test_global_wakefield_weighted(self, ring_with_at_lattice,
+ generate_wakefield):
+ ring_with_at_lattice.optics.local_beta = np.array([0.5, 1])
+ model = ImpedanceModel(ring_with_at_lattice,
+ average_beta=np.array([4, 2]))
+
+ global_wakefield = generate_wakefield(name="global_wake",
+ wake_function_params=[{
+ 'component_type':
+ "xdip",
+ 'function':
+ np.array([1, 2])
+ }, {
+ 'component_type':
+ "ydip",
+ 'function':
+ np.array([1, 2])
+ }],
+ impedance_params=[{
+ 'component_type':
+ "xdip",
+ 'function':
+ np.array([1, 2])
+ }, {
+ 'component_type':
+ "ydip",
+ 'function':
+ np.array([1, 2])
+ }])
+
model.add_global(global_wakefield)
model.compute_sum()
- excepted_x = np.array([4/0.5*1, 4/0.5*2])
- excepted_y = np.array([2/1*1, 2/1*2])
+ excepted_x = np.array([4 / 0.5 * 1, 4 / 0.5 * 2])
+ excepted_y = np.array([2 / 1 * 1, 2 / 1 * 2])
np.testing.assert_allclose(model.sum.Wxdip.data["real"], excepted_x)
np.testing.assert_allclose(model.sum.Zxdip.data["real"], excepted_x)
np.testing.assert_allclose(model.sum.Wydip.data["real"], excepted_y)
np.testing.assert_allclose(model.sum.Zydip.data["real"], excepted_y)
-
+
# Test beta weighting summation
@pytest.mark.parametrize('comp,expected',
[('long', np.array([3, 6])),
- ('xcst', np.array([(2**0.5 + 4**0.5 + 6**0.5)/2, (2**0.5 + 4**0.5 + 6**0.5)/2*2])),
- ('ycst', np.array([(1**0.5 + 2**0.5 + 3**0.5)/1, (1**0.5 + 2**0.5 + 3**0.5)/1*2])),
+ ('xcst',
+ np.array([(2**0.5 + 4**0.5 + 6**0.5) / 2,
+ (2**0.5 + 4**0.5 + 6**0.5) / 2 * 2])),
+ ('ycst',
+ np.array([(1**0.5 + 2**0.5 + 3**0.5) / 1,
+ (1**0.5 + 2**0.5 + 3**0.5) / 1 * 2])),
('xdip', np.array([6, 12])),
('ydip', np.array([6, 12])),
('xquad', np.array([6, 12])),
('yquad', np.array([6, 12]))])
def test_sum_beta(self, demo_ring, generate_wakefield, comp, expected):
-
- wake_in = generate_wakefield(
- wake_function_params=[{'component_type': comp,
- 'function': np.array([1, 2])}],
- impedance_params=[{'component_type': comp,
- 'function': np.array([1, 2])}])
-
- demo_ring.optics.local_beta = np.array([2,1])
+
+ wake_in = generate_wakefield(wake_function_params=[{
+ 'component_type':
+ comp,
+ 'function':
+ np.array([1, 2])
+ }],
+ impedance_params=[{
+ 'component_type':
+ comp,
+ 'function':
+ np.array([1, 2])
+ }])
+
+ demo_ring.optics.local_beta = np.array([2, 1])
model = ImpedanceModel(demo_ring)
model.add(wake_in, 0, name="test")
- beta = np.array([[2,4,6],[1,2,3]])
+ beta = np.array([[2, 4, 6], [1, 2, 3]])
result_wake = model.sum_beta(model.wakefields[0], beta)
wf = getattr(result_wake, f'W{comp}')
z = getattr(result_wake, f'Z{comp}')
assert np.allclose(wf.data["real"], expected)
assert np.allclose(z.data["real"], expected)
-
+
def test_compute_sum_names(self, ring_with_at_lattice, generate_wakefield):
wake1 = generate_wakefield(name="wake1")
wake2 = generate_wakefield(name="wake2")
model = ImpedanceModel(ring_with_at_lattice)
- model.add(wake1, [1,2,3])
- model.add(wake2, [5,6,7])
+ model.add(wake1, [1, 2, 3])
+ model.add(wake2, [5, 6, 7])
model.compute_sum_names()
assert hasattr(model, 'sum_wake1')
assert hasattr(model, 'sum_wake2')
@@ -114,7 +140,8 @@ class TestImpedanceModel:
assert comp in model.sum.components
# Saving the ImpedanceModel to a file should serialize all relevant attributes
- def test_save_impedance_model(self, tmp_path, ring_with_at_lattice, generate_wakefield):
+ def test_save_impedance_model(self, tmp_path, ring_with_at_lattice,
+ generate_wakefield):
model = ImpedanceModel(ring_with_at_lattice)
wakefield = generate_wakefield(name="wake")
model.add(wakefield, [0])
@@ -124,7 +151,8 @@ class TestImpedanceModel:
assert file_path.exists()
# Loading the ImpedanceModel from a file should restore all attributes accurately
- def test_load_impedance_model(self, tmp_path, ring_with_at_lattice, generate_wakefield):
+ def test_load_impedance_model(self, tmp_path, ring_with_at_lattice,
+ generate_wakefield):
model = ImpedanceModel(ring_with_at_lattice)
wakefield = generate_wakefield(name="wake")
global_wake = generate_wakefield(name="global_wake")
@@ -133,60 +161,63 @@ class TestImpedanceModel:
model.compute_sum()
file_path = tmp_path / "impedance_model.pkl"
model.save(file_path)
-
+
new_model = ImpedanceModel(ring_with_at_lattice)
new_model.load(file_path)
-
+
assert "wake" in new_model.names
assert isinstance(new_model.wakefields[0], WakeField)
assert [0] in new_model.positions
assert "global_wake" in new_model.globals_names
assert isinstance(new_model.globals[0], WakeField)
- assert new_model.average_beta.shape == (2,1)
+ assert new_model.average_beta.shape == (2, 1)
# Plotting the contributions of WakeFields should generate a valid plot
- @pytest.mark.parametrize('comp',[('long'),('xdip'),('ydip'),('xquad'),('yquad')])
+ @pytest.mark.parametrize('comp', [('long'), ('xdip'), ('ydip'), ('xquad'),
+ ('yquad')])
def test_plot_area(self, ring_with_at_lattice, generate_wakefield, comp):
model = ImpedanceModel(ring_with_at_lattice)
- wake = generate_wakefield(
- impedance_params=[{'component_type': comp}],
- name="wake"
- )
+ wake = generate_wakefield(impedance_params=[{
+ 'component_type': comp
+ }],
+ name="wake")
model.add(wake, [0])
model.compute_sum()
fig = model.plot_area(Z_type=f"Z{comp}", sigma=30e-12, zoom=True)
assert fig is not None
# Adding a WakeField with a duplicate name should raise a ValueError
- def test_add_duplicate_name_raises_valueerror(self, ring_with_at_lattice, generate_wakefield):
+ def test_add_duplicate_name_raises_valueerror(self, ring_with_at_lattice,
+ generate_wakefield):
model = ImpedanceModel(ring_with_at_lattice)
wakefield = generate_wakefield()
wakefield.name = "duplicate"
-
+
model.add(wakefield, [0], "duplicate")
-
+
with pytest.raises(ValueError):
model.add(wakefield, [1], "duplicate")
# Adding a global WakeField with a duplicate name should raise a ValueError
- def test_add_global_duplicate_name_raises_valueerror(self, ring_with_at_lattice, generate_wakefield):
+ def test_add_global_duplicate_name_raises_valueerror(
+ self, ring_with_at_lattice, generate_wakefield):
model = ImpedanceModel(ring_with_at_lattice)
global_wakefield = generate_wakefield()
global_wakefield.name = "global_duplicate"
-
+
model.add_global(global_wakefield, "global_duplicate")
-
+
with pytest.raises(ValueError):
model.add_global(global_wakefield, "global_duplicate")
# Renaming an attribute should update the internal dictionary correctly
def test_rename_attribute(self, ring_with_at_lattice):
model = ImpedanceModel(ring_with_at_lattice)
-
+
model.some_attribute = 123
-
+
model.rename_attribute("some_attribute", "new_attribute_name")
-
+
assert hasattr(model, "new_attribute_name")
assert not hasattr(model, "some_attribute")
@@ -194,56 +225,62 @@ class TestImpedanceModel:
def test_group_attributes(self, demo_ring, generate_wakefield):
# Setup
model = ImpedanceModel(demo_ring)
-
+
# Create WakeFields with different properties
- wake1 = generate_wakefield(impedance_params=[{'component_type': 'long'},
- {'component_type': 'xdip'}])
- wake2 = generate_wakefield(impedance_params=[{'component_type': 'long'},
- {'component_type': 'xdip'}])
+ wake1 = generate_wakefield(impedance_params=[{
+ 'component_type': 'long'
+ }, {
+ 'component_type': 'xdip'
+ }])
+ wake2 = generate_wakefield(impedance_params=[{
+ 'component_type': 'long'
+ }, {
+ 'component_type': 'xdip'
+ }])
# Add WakeFields to the model
model.add(wake1, positions=[0.0], name='wake1')
model.add(wake2, positions=[1.0], name='wake2')
-
+
# Compute sum to initialize sum_names
model.compute_sum()
-
+
# Group attributes with a common property
model.group_attributes('wake', property_list=['Zlong', 'Zxdip'])
-
+
# Assertions
assert 'wake' in model.sum_names
assert not hasattr(model, 'sum_wake1')
assert not hasattr(model, 'sum_wake2')
-
+
# Check if the grouped attribute has combined properties
grouped_wake = getattr(model, 'wake')
assert hasattr(grouped_wake, 'Zlong')
assert hasattr(grouped_wake, 'Zxdip')
-
+
# Grouping attributes should correctly combine specified properties
# def test_group_attributes_combines_properties(self, demo_ring, generate_wakefield):
# # Setup
# model = ImpedanceModel(demo_ring)
-
+
# # Create WakeFields with different properties
- # wake1 = generate_wakefield(impedance_params=[{'component_type': 'long'},
+ # wake1 = generate_wakefield(impedance_params=[{'component_type': 'long'},
# {'component_type': 'xdip'}])
- # wake2 = generate_wakefield(impedance_params=[{'component_type': 'long'}])
+ # wake2 = generate_wakefield(impedance_params=[{'component_type': 'long'}])
# # Add WakeFields to the model
# model.add(wake1, positions=[0.0], name='wake1')
# model.add(wake2, positions=[1.0], name='wake2')
-
+
# # Compute sum to initialize sum_names
# model.compute_sum()
-
+
# # Group attributes with a common property
# model.group_attributes('wake', property_list=['Zlong', 'Zxdip'])
-
+
# # Assertions
# assert 'wake' in model.sum_names
# assert not hasattr(model, 'sum_wake1')
# assert not hasattr(model, 'sum_wake2')
-
+
# # Check if the grouped attribute has combined properties
# grouped_wake = getattr(model, 'wake')
# assert hasattr(grouped_wake, 'Zlong')
@@ -251,7 +288,10 @@ class TestImpedanceModel:
def test_power_loss_spectrum(self, demo_ring, generate_wakefield):
# Setup
- wakefield = generate_wakefield(impedance_params=[{'component_type': 'long'}])
+ wakefield = generate_wakefield(
+ impedance_params=[{
+ 'component_type': 'long'
+ }])
model = ImpedanceModel(demo_ring)
model.add(wakefield, positions=[0], name='test_wakefield')
model.compute_sum()
@@ -263,7 +303,10 @@ class TestImpedanceModel:
sigma = 1e-9
# Test without max_overlap
- pf0, power_loss = model.power_loss_spectrum(M, bunch_spacing, I, sigma=sigma)
+ pf0, power_loss = model.power_loss_spectrum(M,
+ bunch_spacing,
+ I,
+ sigma=sigma)
# Assertions
assert len(pf0) == len(power_loss)
@@ -273,8 +316,9 @@ class TestImpedanceModel:
def test_energy_loss(self, demo_ring, generate_wakefield):
# Setup
wakefield = generate_wakefield(
- impedance_params=[{'component_type': 'long'}]
- )
+ impedance_params=[{
+ 'component_type': 'long'
+ }])
model = ImpedanceModel(demo_ring)
model.add(wakefield, positions=[0], name='test_wakefield')
model.compute_sum()
@@ -291,4 +335,4 @@ class TestImpedanceModel:
# Verify
assert not summary.empty, "The summary should not be empty"
assert "loss factor (beam) [V/pC]" in summary.columns, "Expected column missing"
- assert "loss factor (bunch) [V/pC]" in summary.columns, "Expected column missing"
\ No newline at end of file
+ assert "loss factor (bunch) [V/pC]" in summary.columns, "Expected column missing"
diff --git a/tests/unit/impedance/test_wakefield.py b/tests/unit/impedance/test_wakefield.py
index fbc7dd08aa98dd114594fea27e3fa8f26b0ec33c..8862ef8a5d25723742c986d6ea059182e522b910 100644
--- a/tests/unit/impedance/test_wakefield.py
+++ b/tests/unit/impedance/test_wakefield.py
@@ -1,20 +1,24 @@
import numpy as np
import pytest
+
pytestmark = pytest.mark.unit
-from mbtrack2 import ComplexData, WakeFunction, Impedance, WakeField
+from mbtrack2 import ComplexData, Impedance, WakeField, WakeFunction
component_list = ComplexData.name_and_coefficients_table().columns.to_list()
+
class TestComplexData:
-
+
@pytest.fixture
def complex_data(self):
- complex_data = ComplexData(variable=np.array([0, 1]), function=np.array([1+1j, 2+2j]))
+ complex_data = ComplexData(variable=np.array([0, 1]),
+ function=np.array([1 + 1j, 2 + 2j]))
return complex_data
# Add two ComplexData objects using 'zero' method and check resulting DataFrame
def test_add_zero_method(self, complex_data):
- complex_data2 = ComplexData(variable=np.array([0, 1]), function=np.array([3+3j, 4+4j]))
+ complex_data2 = ComplexData(variable=np.array([0, 1]),
+ function=np.array([3 + 3j, 4 + 4j]))
result = complex_data + complex_data2
expected_real = np.array([4, 6])
expected_imag = np.array([4, 6])
@@ -31,8 +35,9 @@ class TestComplexData:
# Interpolate ComplexData using __call__ and verify the interpolated values
def test_interpolation_call(self):
- complex_data = ComplexData(variable=np.array([0, 0.25, 0.75, 1]),
- function=np.array([1+1j, 1.25+1.25j, 1.75+1.75j, 2+2j]))
+ complex_data = ComplexData(
+ variable=np.array([0, 0.25, 0.75, 1]),
+ function=np.array([1 + 1j, 1.25 + 1.25j, 1.75 + 1.75j, 2 + 2j]))
interpolated_values = complex_data(np.array([0.5]))
expected_values = np.array([1.5 + 1.5j])
assert np.allclose(interpolated_values, expected_values)
@@ -58,29 +63,37 @@ class TestComplexData:
def test_interpolation_outside_range(self, complex_data):
with pytest.raises(ValueError):
complex_data(np.array([-0.5, 1.5]))
-
+
# Interpolate ComplexData with values outside the index range
def test_wrong_component_type(self, complex_data):
with pytest.raises(KeyError):
complex_data.component_type = "wrong"
+
@pytest.fixture(params=component_list)
def wf_param(request):
- wf = WakeFunction(variable=np.array([0, 1]), function=np.array([1, 2]), component_type=request.param)
+ wf = WakeFunction(variable=np.array([0, 1]),
+ function=np.array([1, 2]),
+ component_type=request.param)
return wf
+
@pytest.fixture
def generate_wf():
- def generate(variable=np.array([0, 1]),
- function=np.array([1, 2]),
- component_type='long',
- ):
+
+ def generate(
+ variable=np.array([0, 1]),
+ function=np.array([1, 2]),
+ component_type='long',
+ ):
wf = WakeFunction(variable=variable,
- function=function,
+ function=function,
component_type=component_type)
return wf
+
return generate
+
class TestWakeFunction:
# Add two WakeFunction objects with the same component type and verify result
@@ -97,26 +110,28 @@ class TestWakeFunction:
result = wf * 2
assert np.allclose(result.data['real'], [2, 4])
assert result.component_type == 'long'
-
+
# Add WakeFunction objects with different component types and check for warnings
def test_add_different_component_types_warning(self, generate_wf):
wf1 = generate_wf(component_type='long')
wf2 = generate_wf(component_type='xdip')
- with pytest.warns(UserWarning, match="do not have the same coordinates or plane or type"):
+ with pytest.warns(
+ UserWarning,
+ match="do not have the same coordinates or plane or type"):
result = wf1 + wf2
assert result.data.equals(wf1.data)
# Convert WakeFunction to Impedance and verify the transformation
def test_convert_to_impedance(self, generate_wf):
- for component_type in ["long","xdip","ydip"]:
+ for component_type in ["long", "xdip", "ydip"]:
wf = generate_wf(component_type=component_type)
imp = wf.to_impedance(freq_lim=10)
assert imp.component_type == component_type
assert isinstance(imp, Impedance)
-
+
# Test deconvolution
def test_deconvolution(self, generate_wf):
- for component_type in ["long","xdip","ydip"]:
+ for component_type in ["long", "xdip", "ydip"]:
wf = generate_wf(component_type=component_type)
deconv_wf = wf.deconvolution(freq_lim=1e9, sigma=0.01, mu=0)
assert isinstance(deconv_wf, WakeFunction)
@@ -130,30 +145,38 @@ class TestWakeFunction:
# Plot the WakeFunction.loss_factor
def test_loss_factor(self, wf_param):
assert wf_param.loss_factor(1) > 0
-
+
+
@pytest.fixture(params=component_list)
def imp_param(request):
- imp = Impedance(variable=np.array([1, 2]), function=np.array([3+4j, 5+6j]), component_type=request.param)
+ imp = Impedance(variable=np.array([1, 2]),
+ function=np.array([3 + 4j, 5 + 6j]),
+ component_type=request.param)
return imp
+
@pytest.fixture
def generate_imp():
- def generate(variable=np.array([1, 2]),
- function=np.array([3+4j, 5+6j]),
- component_type='long',
- ):
+
+ def generate(
+ variable=np.array([1, 2]),
+ function=np.array([3 + 4j, 5 + 6j]),
+ component_type='long',
+ ):
imp = Impedance(variable=variable,
- function=function,
+ function=function,
component_type=component_type)
return imp
+
return generate
-
+
+
class TestImpedance:
# Impedance objects can be added using the add method with matching component types
def test_add_matching_component_types(self, generate_imp):
imp1 = generate_imp()
- imp2 = generate_imp(function=np.array([1+1j, 1+1j]))
+ imp2 = generate_imp(function=np.array([1 + 1j, 1 + 1j]))
result = imp1 + imp2
expected_real = np.array([4, 6])
expected_imag = np.array([5, 7])
@@ -173,13 +196,13 @@ class TestImpedance:
# The loss_factor method computes the correct loss factor for a given sigma
def test_loss_factor_computation(self, generate_imp):
- for component_type in ["long","xdip","ydip","xquad","yquad"]:
+ for component_type in ["long", "xdip", "ydip", "xquad", "yquad"]:
imp = generate_imp(component_type=component_type)
assert imp.loss_factor(1) > 0
# The to_wakefunction method correctly transforms impedance data to a WakeFunction object
def test_to_wakefunction_transformation(self, generate_imp):
- for component_type in ["long","xdip","ydip"]:
+ for component_type in ["long", "xdip", "ydip"]:
imp = generate_imp(component_type=component_type)
wf = imp.to_wakefunction()
assert isinstance(wf, WakeFunction)
@@ -190,8 +213,10 @@ class TestImpedance:
imp_param.plot()
assert True
+
@pytest.fixture
def generate_wakefield(generate_imp, generate_wf):
+
def generate(impedance_params=None, wake_function_params=None, name=None):
structure_list = []
if impedance_params:
@@ -201,28 +226,40 @@ def generate_wakefield(generate_imp, generate_wf):
for params in wake_function_params:
structure_list.append(generate_wf(**params))
return WakeField(structure_list=structure_list, name=name)
+
return generate
+
@pytest.fixture(params=component_list)
def wakefield_param(generate_wakefield, request):
- wake = generate_wakefield(
- impedance_params=[{'component_type': request.param}],
- wake_function_params=[{'component_type': request.param}]
- )
+ wake = generate_wakefield(impedance_params=[{
+ 'component_type': request.param
+ }],
+ wake_function_params=[{
+ 'component_type':
+ request.param
+ }])
return wake
+
class TestWakeField:
- def test_initialize_with_impedance_wakefunction_list(self, generate_wakefield):
+ def test_initialize_with_impedance_wakefunction_list(
+ self, generate_wakefield):
for component_type in component_list:
- wakefield = generate_wakefield(
- impedance_params=[{'component_type': component_type}],
- wake_function_params=[{'component_type': component_type}]
- )
+ wakefield = generate_wakefield(impedance_params=[{
+ 'component_type':
+ component_type
+ }],
+ wake_function_params=[{
+ 'component_type':
+ component_type
+ }])
assert f'Z{component_type}' in wakefield.components
assert f'W{component_type}' in wakefield.components
- def test_append_to_model(self, generate_imp, generate_wakefield, generate_wf):
+ def test_append_to_model(self, generate_imp, generate_wakefield,
+ generate_wf):
wakefield = generate_wakefield()
for component_type in component_list:
imp = generate_imp(component_type=component_type)
@@ -234,53 +271,56 @@ class TestWakeField:
def test_save_and_load_wakefield(self, generate_wakefield, tmp_path):
wakefield = generate_wakefield(
- impedance_params=[{'component_type': 'long'}]
- )
+ impedance_params=[{
+ 'component_type': 'long'
+ }])
file_path = tmp_path / "wakefield.pkl"
wakefield.save(file_path)
loaded_wakefield = WakeField.load(file_path)
assert 'Zlong' in loaded_wakefield.components
- @pytest.mark.parametrize('comp,expected',
- [('long', np.array([4, 6])),
- ('xcst', np.array([3 + np.sqrt(2), 4 + 2*np.sqrt(2)])),
- ('ycst', np.array([3 + np.sqrt(2), 4 + 2*np.sqrt(2)])),
- ('xdip', np.array([5, 8])),
- ('ydip', np.array([5, 8])),
- ('xquad', np.array([5, 8])),
- ('yquad', np.array([5, 8]))]
- )
- def test_add_wakefields_with_beta_functions(self,
- generate_wakefield,
- comp,
+ @pytest.mark.parametrize(
+ 'comp,expected',
+ [('long', np.array([4, 6])),
+ ('xcst', np.array([3 + np.sqrt(2), 4 + 2 * np.sqrt(2)])),
+ ('ycst', np.array([3 + np.sqrt(2), 4 + 2 * np.sqrt(2)])),
+ ('xdip', np.array([5, 8])), ('ydip', np.array([5, 8])),
+ ('xquad', np.array([5, 8])), ('yquad', np.array([5, 8]))])
+ def test_add_wakefields_with_beta_functions(self, generate_wakefield, comp,
expected):
wake1 = generate_wakefield(
- wake_function_params=[{'component_type': comp,
- 'function': np.array([1, 2])}]
- )
+ wake_function_params=[{
+ 'component_type': comp,
+ 'function': np.array([1, 2])
+ }])
wake2 = generate_wakefield(
- wake_function_params=[{'component_type': comp,
- 'function': np.array([3, 4])}]
- )
+ wake_function_params=[{
+ 'component_type': comp,
+ 'function': np.array([3, 4])
+ }])
beta1 = [2, 2]
beta2 = [1, 1]
-
+
result_wake = WakeField.add_wakefields(wake1, beta1, wake2, beta2)
wf = getattr(result_wake, f'W{comp}')
assert np.allclose(wf.data["real"], expected)
-
- result_wake = WakeField.add_several_wakefields([wake1, wake2], np.array([beta1, beta2]))
+
+ result_wake = WakeField.add_several_wakefields([wake1, wake2],
+ np.array([beta1,
+ beta2]))
wf = getattr(result_wake, f'W{comp}')
assert np.allclose(wf.data["real"], expected)
def test_add_duplicate_components_raises_error(self, generate_wakefield):
wakefield = generate_wakefield(
- impedance_params=[{'component_type': 'long'}]
- )
+ impedance_params=[{
+ 'component_type': 'long'
+ }])
with pytest.raises(ValueError):
wakefield.append_to_model(wakefield.Zlong)
- def test_drop_non_existent_component_raises_error(self, generate_wakefield):
+ def test_drop_non_existent_component_raises_error(self,
+ generate_wakefield):
wakefield = generate_wakefield()
with pytest.raises(AttributeError):
wakefield.drop('Znonexistent')
@@ -288,35 +328,47 @@ class TestWakeField:
def test_append_to_model_with_invalid_component(self, generate_wakefield):
wakefield = generate_wakefield()
invalid_component = "InvalidComponent"
- with pytest.raises(ValueError, match="is not an Impedance nor a WakeFunction."):
+ with pytest.raises(ValueError,
+ match="is not an Impedance nor a WakeFunction."):
wakefield.append_to_model(invalid_component)
def test_drop_method_input_types(self, generate_wakefield):
- wakefield = generate_wakefield(
- impedance_params=[{'component_type': 'long'}],
- wake_function_params=[{'component_type': 'xdip'}]
- )
-
+ wakefield = generate_wakefield(impedance_params=[{
+ 'component_type':
+ 'long'
+ }],
+ wake_function_params=[{
+ 'component_type':
+ 'xdip'
+ }])
+
wakefield.drop('Zlong')
assert 'Zlong' not in wakefield.components
-
+
wakefield.drop(['Wxdip'])
assert 'Wxdip' not in wakefield.components
-
- wakefield = generate_wakefield(
- impedance_params=[{'component_type': 'long'}],
- wake_function_params=[{'component_type': 'xdip'}]
- )
+
+ wakefield = generate_wakefield(impedance_params=[{
+ 'component_type':
+ 'long'
+ }],
+ wake_function_params=[{
+ 'component_type':
+ 'xdip'
+ }])
wakefield.drop('Z')
assert 'Zlong' not in wakefield.components
-
- wakefield = generate_wakefield(
- impedance_params=[{'component_type': 'long'}],
- wake_function_params=[{'component_type': 'xdip'}]
- )
+
+ wakefield = generate_wakefield(impedance_params=[{
+ 'component_type':
+ 'long'
+ }],
+ wake_function_params=[{
+ 'component_type':
+ 'xdip'
+ }])
wakefield.drop('W')
assert 'Wxdip' not in wakefield.components
-
+
with pytest.raises(TypeError):
wakefield.drop(123)
-
\ No newline at end of file
diff --git a/tests/unit/tracking/test_aperture.py b/tests/unit/tracking/test_aperture.py
index fd8a1a2589efcaeabaca2987d70a5726a5b757a2..91a5833b383e2ca99d10543a89ca956f6a74ad60 100644
--- a/tests/unit/tracking/test_aperture.py
+++ b/tests/unit/tracking/test_aperture.py
@@ -1,9 +1,13 @@
import pytest
+
pytestmark = pytest.mark.unit
-from mbtrack2 import (CircularAperture,
- ElipticalAperture,
- RectangularAperture,
- LongitudinalAperture)
+from mbtrack2 import (
+ CircularAperture,
+ ElipticalAperture,
+ LongitudinalAperture,
+ RectangularAperture,
+)
+
class TestCircularAperture:
@@ -16,10 +20,8 @@ class TestCircularAperture:
assert all(small_bunch.alive)
# Track a bunch with particles exactly on the radius boundary
- @pytest.mark.parametrize("x, y", [(1.0, 0.0),
- (0.0, 1.0),
- (-1.0, 0.0),
- (0.0, -1.0)])
+ @pytest.mark.parametrize("x, y", [(1.0, 0.0), (0.0, 1.0), (-1.0, 0.0),
+ (0.0, -1.0)])
def test_track_particles_on_boundary(self, small_bunch, x, y):
aperture = CircularAperture(radius=1.0)
small_bunch["x"] = x
@@ -28,9 +30,7 @@ class TestCircularAperture:
assert all(small_bunch.alive)
# Track a bunch with all particles outside the radius
- @pytest.mark.parametrize("x, y", [(0.75, 0.75),
- (0.75, -0.75),
- (1.1, 0.0),
+ @pytest.mark.parametrize("x, y", [(0.75, 0.75), (0.75, -0.75), (1.1, 0.0),
(0.0, 1.1)])
def test_track_all_particles_outside_radius(self, small_bunch, x, y):
aperture = CircularAperture(radius=1.0)
@@ -46,6 +46,7 @@ class TestCircularAperture:
aperture.track(small_bunch)
assert len(small_bunch) == 0
+
class TestElipticalAperture:
# Track a bunch where all particles are within the elliptical aperture
@@ -57,9 +58,7 @@ class TestElipticalAperture:
assert all(small_bunch.alive)
# Track a bunch where some particles are outside the elliptical aperture
- @pytest.mark.parametrize("x, y", [(0.5, 1.75),
- (0.5, -1.75),
- (1.1, 0.0),
+ @pytest.mark.parametrize("x, y", [(0.5, 1.75), (0.5, -1.75), (1.1, 0.0),
(0.0, 2.1)])
def test_track_particles_outside_aperture(self, small_bunch, x, y):
aperture = ElipticalAperture(X_radius=1.0, Y_radius=2.0)
@@ -67,19 +66,17 @@ class TestElipticalAperture:
small_bunch["y"] = y
aperture.track(small_bunch)
assert not all(small_bunch.alive)
-
+
# Track a bunch with no particles
def test_track_no_particles(self, small_bunch):
aperture = ElipticalAperture(X_radius=1.0, Y_radius=2.0)
small_bunch.alive[:] = False
aperture.track(small_bunch)
assert len(small_bunch) == 0
-
+
# Track a bunch with particles exactly on the radius boundary
- @pytest.mark.parametrize("x, y", [(1.0, 0.0),
- (0.0, 2.0),
- (-1.0, 0.0),
- (0.0, -2.0)])
+ @pytest.mark.parametrize("x, y", [(1.0, 0.0), (0.0, 2.0), (-1.0, 0.0),
+ (0.0, -2.0)])
def test_track_particles_on_boundary(self, small_bunch, x, y):
aperture = ElipticalAperture(X_radius=1.0, Y_radius=2.0)
small_bunch["x"] = x
@@ -87,6 +84,7 @@ class TestElipticalAperture:
aperture.track(small_bunch)
assert all(small_bunch.alive)
+
class TestRectangularAperture:
# Track method correctly identifies particles within the rectangular aperture
@@ -98,9 +96,7 @@ class TestRectangularAperture:
assert all(small_bunch.alive)
# Particles outside the defined apertures are marked as 'lost'
- @pytest.mark.parametrize("x, y", [(1.1, 1.1),
- (-1.1, -1.1),
- (1.1, 0.0),
+ @pytest.mark.parametrize("x, y", [(1.1, 1.1), (-1.1, -1.1), (1.1, 0.0),
(0.0, 1.1)])
def test_particles_outside_aperture_lost(self, small_bunch, x, y):
aperture = RectangularAperture(X_right=1.0, Y_top=1.0)
@@ -115,40 +111,41 @@ class TestRectangularAperture:
small_bunch.alive[:] = False
aperture.track(small_bunch)
assert len(small_bunch) == 0
-
+
# Manages particles exactly on the aperture boundaries
- @pytest.mark.parametrize("x, y", [(1.0, 0.0),
- (0.0, 1.0),
- (-1.0, 0.0),
- (0.0, -1.0),
- (1.0, 1.0)])
+ @pytest.mark.parametrize("x, y", [(1.0, 0.0), (0.0, 1.0), (-1.0, 0.0),
+ (0.0, -1.0), (1.0, 1.0)])
def test_particles_on_boundary(self, small_bunch, x, y):
aperture = RectangularAperture(X_right=1.0, Y_top=1.0)
small_bunch.particles["x"] = x
small_bunch.particles["y"] = y
aperture.track(small_bunch)
assert all(small_bunch.alive)
-
- @pytest.mark.parametrize("x, y", [(1.1, 1.1),
- (-1.1, -1.1),
- (0.0, 0.0),
+
+ @pytest.mark.parametrize("x, y", [(1.1, 1.1), (-1.1, -1.1), (0.0, 0.0),
(0.5, -0.6)])
def test_non_default_rect_loss(self, small_bunch, x, y):
- aperture = RectangularAperture(X_right=1.0, Y_top=1.0, X_left=0.1, Y_bottom=-0.5)
+ aperture = RectangularAperture(X_right=1.0,
+ Y_top=1.0,
+ X_left=0.1,
+ Y_bottom=-0.5)
small_bunch["x"] = x
small_bunch["y"] = y
aperture.track(small_bunch)
assert not any(small_bunch.alive)
- @pytest.mark.parametrize("x, y", [(0.5, 0.0),
- (0.2, -0.3)])
+ @pytest.mark.parametrize("x, y", [(0.5, 0.0), (0.2, -0.3)])
def test_non_default_rect(self, small_bunch, x, y):
- aperture = RectangularAperture(X_right=1.0, Y_top=1.0, X_left=0.1, Y_bottom=-0.5)
+ aperture = RectangularAperture(X_right=1.0,
+ Y_top=1.0,
+ X_left=0.1,
+ Y_bottom=-0.5)
small_bunch["x"] = x
small_bunch["y"] = y
aperture.track(small_bunch)
assert all(small_bunch.alive)
-
+
+
class TestLongitudinalAperture:
# Track a Bunch object with particles within the longitudinal bounds
@@ -159,7 +156,7 @@ class TestLongitudinalAperture:
assert all(small_bunch.alive)
# Track a Bunch object with particles exactly on the boundary values
- @pytest.mark.parametrize("tau",[(-1.0),(1.0)])
+ @pytest.mark.parametrize("tau", [(-1.0), (1.0)])
def test_track_bunch_on_boundary(self, small_bunch, tau):
aperture = LongitudinalAperture(tau_up=1.0)
small_bunch["tau"] = tau
@@ -174,9 +171,9 @@ class TestLongitudinalAperture:
assert len(small_bunch) == 0
# Track a Bunch object with all particles outside the longitudinal bounds
- @pytest.mark.parametrize("tau", [(1.1),(-0.6)])
+ @pytest.mark.parametrize("tau", [(1.1), (-0.6)])
def test_track_bunch_outside_bounds(self, small_bunch, tau):
aperture = LongitudinalAperture(tau_up=1.0, tau_low=-0.5)
small_bunch["tau"] = tau
aperture.track(small_bunch)
- assert not any(small_bunch.alive)
\ No newline at end of file
+ assert not any(small_bunch.alive)
diff --git a/tests/unit/tracking/test_beam_ion_effects.py b/tests/unit/tracking/test_beam_ion_effects.py
index ca42a7c60bc0189a98db6005b569542e7283a944..fda29739d775e4e037a28abed5e0f5d91e6e7d7b 100644
--- a/tests/unit/tracking/test_beam_ion_effects.py
+++ b/tests/unit/tracking/test_beam_ion_effects.py
@@ -1,13 +1,18 @@
-from mbtrack2 import IonMonitor, IonParticles, IonAperture, BeamIonElement
-from utility_test_functions import assert_attr_changed
import os
+
import h5py as hp
import numpy as np
import pytest
+from utility_test_functions import assert_attr_changed
+
+from mbtrack2 import BeamIonElement, IonAperture, IonMonitor, IonParticles
+
pytestmark = pytest.mark.unit
+
@pytest.fixture
def generate_ion_particles(demo_ring):
+
def generate(mp_number=100,
ion_element_length=1.0,
ring=demo_ring,
@@ -19,36 +24,55 @@ def generate_ion_particles(demo_ring):
track_alive=track_alive,
alive=alive)
return ions
+
return generate
+
class TestIonParticles:
# Generate particle distribution using electron bunch parameters via generate_as_a_distribution()
- def test_generate_as_distribution(self, generate_ion_particles, large_bunch):
+ def test_generate_as_distribution(self, generate_ion_particles,
+ large_bunch):
mp_number = 1e5
charge = 1e-9
ions = generate_ion_particles(mp_number)
ions.generate_as_a_distribution(large_bunch, charge)
-
- assert np.isclose(ions["x"].mean(), large_bunch["x"].mean(), rtol=0.1, atol=1e-5)
- assert np.isclose(ions["x"].std(), large_bunch["x"].std(), rtol=0.1, atol=1e-5)
- assert np.isclose(ions["y"].mean(), large_bunch["y"].mean(), rtol=0.1, atol=1e-5)
- assert np.isclose(ions["y"].std(), large_bunch["y"].std(), rtol=0.1, atol=1e-5)
+
+ assert np.isclose(ions["x"].mean(),
+ large_bunch["x"].mean(),
+ rtol=0.1,
+ atol=1e-5)
+ assert np.isclose(ions["x"].std(),
+ large_bunch["x"].std(),
+ rtol=0.1,
+ atol=1e-5)
+ assert np.isclose(ions["y"].mean(),
+ large_bunch["y"].mean(),
+ rtol=0.1,
+ atol=1e-5)
+ assert np.isclose(ions["y"].std(),
+ large_bunch["y"].std(),
+ rtol=0.1,
+ atol=1e-5)
assert np.all(ions["xp"] == 0)
assert np.all(ions["yp"] == 0)
assert np.all(ions["delta"] == 0)
assert np.all(ions["tau"] >= -ions.ion_element_length)
assert np.all(ions["tau"] <= ions.ion_element_length)
assert np.isclose(ions.charge, charge)
- assert np.isclose(ions["charge"].mean(), charge/mp_number, rtol=0.01, atol=1e-9)
+ assert np.isclose(ions["charge"].mean(),
+ charge / mp_number,
+ rtol=0.01,
+ atol=1e-9)
# Generate random particle samples from electron bunch via generate_from_random_samples()
- def test_generate_from_random_samples(self, generate_ion_particles, large_bunch):
+ def test_generate_from_random_samples(self, generate_ion_particles,
+ large_bunch):
mp_number = 1e5
charge = 1e-9
ions = generate_ion_particles(mp_number)
ions.generate_from_random_samples(large_bunch, charge)
-
+
assert np.all(np.isin(ions["x"], large_bunch["x"]))
assert np.all(np.isin(ions["y"], large_bunch["y"]))
assert np.all(ions["xp"] == 0)
@@ -57,51 +81,64 @@ class TestIonParticles:
assert np.all(ions["tau"] >= -ions.ion_element_length)
assert np.all(ions["tau"] <= ions.ion_element_length)
assert np.isclose(ions.charge, charge)
- assert np.isclose(ions["charge"].mean(), charge/mp_number, rtol=0.01, atol=1e-9)
-
+ assert np.isclose(ions["charge"].mean(),
+ charge / mp_number,
+ rtol=0.01,
+ atol=1e-9)
+
# Add two IonParticles instances together and verify combined particle arrays
def test_add_ion_particles(self, generate_ion_particles):
ions1 = generate_ion_particles(mp_number=100)
ions2 = generate_ion_particles(mp_number=50)
-
+
combined = ions1 + ions2
assert combined.mp_number == 150
- assert all(combined[coord].shape == (150,) for coord in ["x","xp","y","yp","tau","delta","charge"])
+ assert all(
+ combined[coord].shape == (150, )
+ for coord in ["x", "xp", "y", "yp", "tau", "delta", "charge"])
assert np.all(combined.alive == True)
# Initialize with alive=False and verify all particles marked as dead
def test_init_dead_particles(self, generate_ion_particles):
ions = generate_ion_particles(alive=False)
assert np.all(ions.alive == False)
- assert all(ions[coord].shape == (1,) for coord in ["x","xp","y","yp","tau","delta","charge"])
+ assert all(
+ ions[coord].shape == (1, )
+ for coord in ["x", "xp", "y", "yp", "tau", "delta", "charge"])
# Generate distributions with electron bunch containing no particles
- def test_generate_from_empty_bunch(self, generate_ion_particles, small_bunch):
+ def test_generate_from_empty_bunch(self, generate_ion_particles,
+ small_bunch):
small_bunch.alive[:] = False
ions = generate_ion_particles()
-
+
with pytest.raises(ValueError):
ions.generate_as_a_distribution(small_bunch, 1e-9)
with pytest.raises(ValueError):
ions.generate_from_random_samples(small_bunch, 1e-9)
-
+
+
@pytest.fixture
def generate_ion_monitor(tmp_path):
- def generate(save_every=1,
- buffer_size=10,
- total_size=10,
+
+ def generate(save_every=1,
+ buffer_size=10,
+ total_size=10,
file_name=tmp_path / "test_monitor.hdf5"):
monitor = IonMonitor(save_every=save_every,
- buffer_size=buffer_size,
- total_size=total_size,
- file_name=file_name)
+ buffer_size=buffer_size,
+ total_size=total_size,
+ file_name=file_name)
return monitor
+
return generate
+
class TestIonMonitor:
# Monitor initialization with valid parameters creates HDF5 file and sets up data structures
- def test_monitor_init_creates_valid_structures(self, generate_ion_monitor, tmp_path):
+ def test_monitor_init_creates_valid_structures(self, generate_ion_monitor,
+ tmp_path):
monitor = generate_ion_monitor()
assert monitor.file is not None
assert monitor.buffer_size == 10
@@ -112,36 +149,40 @@ class TestIonMonitor:
assert monitor.track_count == 0
# Buffer writes to file when full and resets counter
- def test_buffer_writes_when_full(self, generate_ion_monitor, generate_ion_particles):
+ def test_buffer_writes_when_full(self, generate_ion_monitor,
+ generate_ion_particles):
monitor = generate_ion_monitor(buffer_size=2, total_size=4)
ions = generate_ion_particles()
-
+
for _ in range(2):
monitor.track(ions)
-
+
assert monitor.buffer_count == 0
assert monitor.write_count == 1
# Data structures are properly initialized with correct shapes and types
def test_data_structures_initialization(self, generate_ion_monitor):
monitor = generate_ion_monitor()
-
+
assert monitor.mean.shape == (4, 10)
assert monitor.std.shape == (4, 10)
- assert monitor.charge.shape == (10,)
- assert monitor.time.shape == (10,)
+ assert monitor.charge.shape == (10, )
+ assert monitor.time.shape == (10, )
assert monitor.time.dtype == int
# Initialize monitor with total_size not divisible by buffer_size
def test_invalid_total_size_raises_error(self, generate_ion_monitor):
- with pytest.raises(ValueError, match="total_size must be divisible by buffer_size"):
+ with pytest.raises(
+ ValueError,
+ match="total_size must be divisible by buffer_size"):
generate_ion_monitor(buffer_size=3, total_size=10)
# Initialize with invalid/missing file_name
def test_invalid_filename_handling(self, generate_ion_monitor):
with pytest.raises(OSError):
generate_ion_monitor(file_name="/invalid/path/file.hdf5")
-
+
+
class TestElipticalAperture:
# Track a bunch where all particles are within the elliptical aperture
@@ -149,131 +190,147 @@ class TestElipticalAperture:
ions = generate_ion_particles()
mp = len(ions)
aperture = IonAperture(X_radius=1.0, Y_radius=2.0)
- ions["x"] = np.ones_like(ions["x"])*0.5
- ions["y"] = np.ones_like(ions["y"])*1.0
+ ions["x"] = np.ones_like(ions["x"]) * 0.5
+ ions["y"] = np.ones_like(ions["y"]) * 1.0
aperture.track(ions)
assert all(ions.alive)
assert mp == len(ions)
# Track a bunch where some particles are outside the elliptical aperture
- @pytest.mark.parametrize("x, y", [(0.5, 1.75),
- (0.5, -1.75),
- (1.1, 0.0),
+ @pytest.mark.parametrize("x, y", [(0.5, 1.75), (0.5, -1.75), (1.1, 0.0),
(0.0, 2.1)])
- def test_track_particles_outside_aperture(self, generate_ion_particles, x, y):
+ def test_track_particles_outside_aperture(self, generate_ion_particles, x,
+ y):
ions = generate_ion_particles()
mp = len(ions)
aperture = IonAperture(X_radius=1.0, Y_radius=2.0)
- ions["x"] = np.ones_like(ions["x"])*x
- ions["y"] = np.ones_like(ions["y"])*y
+ ions["x"] = np.ones_like(ions["x"]) * x
+ ions["y"] = np.ones_like(ions["y"]) * y
aperture.track(ions)
assert mp != len(ions)
-
+
# Track a bunch with no particles
def test_track_no_particles(self, generate_ion_particles):
ions = generate_ion_particles()
- ions["x"] = np.ones_like(ions["x"])*100
- ions["y"] = np.ones_like(ions["y"])*100
+ ions["x"] = np.ones_like(ions["x"]) * 100
+ ions["y"] = np.ones_like(ions["y"]) * 100
aperture = IonAperture(X_radius=1.0, Y_radius=2.0)
aperture.track(ions)
assert len(ions) == 0
-
+
aperture.track(ions)
assert True
-
+
# Track a bunch with particles exactly on the radius boundary
- @pytest.mark.parametrize("x, y", [(1.0, 0.0),
- (0.0, 2.0),
- (-1.0, 0.0),
- (0.0, -2.0)])
+ @pytest.mark.parametrize("x, y", [(1.0, 0.0), (0.0, 2.0), (-1.0, 0.0),
+ (0.0, -2.0)])
def test_track_particles_on_boundary(self, generate_ion_particles, x, y):
ions = generate_ion_particles()
mp = len(ions)
aperture = IonAperture(X_radius=1.0, Y_radius=2.0)
- ions["x"] = np.ones_like(ions["x"])*x
- ions["y"] = np.ones_like(ions["y"])*y
+ ions["x"] = np.ones_like(ions["x"]) * x
+ ions["y"] = np.ones_like(ions["y"]) * y
aperture.track(ions)
assert all(ions.alive)
assert mp == len(ions)
-
+
+
@pytest.fixture
def generate_beam_ion(demo_ring):
- def generate(
- ion_mass=1.67e-27,
- ion_charge=1.6e-19,
- ionization_cross_section=1e-22,
- residual_gas_density=1e50,
- ring=demo_ring,
- ion_field_model="strong",
- electron_field_model="strong",
- ion_element_length=demo_ring.L,
- n_ion_macroparticles_per_bunch=30,
- generate_method='samples'):
-
+
+ def generate(ion_mass=1.67e-27,
+ ion_charge=1.6e-19,
+ ionization_cross_section=1e-22,
+ residual_gas_density=1e50,
+ ring=demo_ring,
+ ion_field_model="strong",
+ electron_field_model="strong",
+ ion_element_length=demo_ring.L,
+ n_ion_macroparticles_per_bunch=30,
+ generate_method='samples'):
+
beam_ion = BeamIonElement(
ion_mass=ion_mass,
- ion_charge=ion_charge,
+ ion_charge=ion_charge,
ionization_cross_section=ionization_cross_section,
residual_gas_density=residual_gas_density,
ring=ring,
ion_field_model=ion_field_model,
- electron_field_model=electron_field_model,
+ electron_field_model=electron_field_model,
ion_element_length=ion_element_length,
n_ion_macroparticles_per_bunch=n_ion_macroparticles_per_bunch,
generate_method=generate_method)
return beam_ion
+
return generate
+
class TestBeamIonElement:
@pytest.mark.parametrize('ion_field_model, electron_field_model',
- [('weak','weak'), ('weak','strong'),
- ('strong','weak'), ('strong', 'strong')])
- def test_track_bunch(self, generate_beam_ion, small_bunch, ion_field_model, electron_field_model):
- beam_ion = generate_beam_ion(ion_field_model=ion_field_model, electron_field_model=electron_field_model)
- assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
-
+ [('weak', 'weak'), ('weak', 'strong'),
+ ('strong', 'weak'), ('strong', 'strong')])
+ def test_track_bunch(self, generate_beam_ion, small_bunch, ion_field_model,
+ electron_field_model):
+ beam_ion = generate_beam_ion(ion_field_model=ion_field_model,
+ electron_field_model=electron_field_model)
+ assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp", "yp"])
+
@pytest.mark.parametrize('ion_field_model, electron_field_model',
- [('weak','weak'), ('weak','strong'),
- ('strong','weak'), ('strong', 'strong')])
- def test_track_bunch_partially_lost(self, generate_beam_ion, small_bunch, ion_field_model, electron_field_model):
- beam_ion = generate_beam_ion(ion_field_model=ion_field_model, electron_field_model=electron_field_model)
- assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+ [('weak', 'weak'), ('weak', 'strong'),
+ ('strong', 'weak'), ('strong', 'strong')])
+ def test_track_bunch_partially_lost(self, generate_beam_ion, small_bunch,
+ ion_field_model, electron_field_model):
+ beam_ion = generate_beam_ion(ion_field_model=ion_field_model,
+ electron_field_model=electron_field_model)
+ assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp", "yp"])
charge_gen = beam_ion.ion_beam.charge
-
+
# loose half of electron bunch
small_bunch.alive[0:5] = False
- assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
-
- assert np.isclose(beam_ion.ion_beam.charge, charge_gen*1.5)
+ assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp", "yp"])
+
+ assert np.isclose(beam_ion.ion_beam.charge, charge_gen * 1.5)
@pytest.mark.parametrize('ion_field_model, electron_field_model',
- [('weak','weak'), ('weak','strong'),
- ('strong','weak'), ('strong', 'strong')])
- def test_track_beam(self, generate_beam_ion, beam_uniform, ion_field_model, electron_field_model):
- beam_ion = generate_beam_ion(ion_field_model=ion_field_model, electron_field_model=electron_field_model)
- assert_attr_changed(beam_ion, beam_uniform, attrs_changed=["xp","yp"])
+ [('weak', 'weak'), ('weak', 'strong'),
+ ('strong', 'weak'), ('strong', 'strong')])
+ def test_track_beam(self, generate_beam_ion, beam_uniform, ion_field_model,
+ electron_field_model):
+ beam_ion = generate_beam_ion(ion_field_model=ion_field_model,
+ electron_field_model=electron_field_model)
+ assert_attr_changed(beam_ion, beam_uniform, attrs_changed=["xp", "yp"])
@pytest.mark.parametrize('ion_field_model, electron_field_model',
- [('weak','weak'), ('weak','strong'),
- ('strong','weak'), ('strong', 'strong')])
- def test_track_beam_non_uniform(self, generate_beam_ion, beam_non_uniform, ion_field_model, electron_field_model):
- beam_ion = generate_beam_ion(ion_field_model=ion_field_model, electron_field_model=electron_field_model)
- assert_attr_changed(beam_ion, beam_non_uniform, attrs_changed=["xp","yp"])
+ [('weak', 'weak'), ('weak', 'strong'),
+ ('strong', 'weak'), ('strong', 'strong')])
+ def test_track_beam_non_uniform(self, generate_beam_ion, beam_non_uniform,
+ ion_field_model, electron_field_model):
+ beam_ion = generate_beam_ion(ion_field_model=ion_field_model,
+ electron_field_model=electron_field_model)
+ assert_attr_changed(beam_ion,
+ beam_non_uniform,
+ attrs_changed=["xp", "yp"])
# Ion generation creates expected number of macroparticles with proper distribution
- @pytest.mark.parametrize('generate_method', [('samples'),('distribution')])
- def test_ion_generation(self, generate_beam_ion, large_bunch, generate_method):
+ @pytest.mark.parametrize('generate_method', [('samples'),
+ ('distribution')])
+ def test_ion_generation(self, generate_beam_ion, large_bunch,
+ generate_method):
n_ions = 1e5
large_bunch["x"] += 1
large_bunch["y"] += 1
beam_ion = generate_beam_ion(n_ion_macroparticles_per_bunch=n_ions,
generate_method=generate_method)
beam_ion.generate_new_ions(large_bunch)
-
+
assert len(beam_ion.ion_beam["x"]) == n_ions + 1
- assert np.isclose(beam_ion.ion_beam["x"].mean(), large_bunch["x"].mean(), rtol=0.1)
- assert np.isclose(beam_ion.ion_beam["y"].mean(), large_bunch["y"].mean(), rtol=0.1)
+ assert np.isclose(beam_ion.ion_beam["x"].mean(),
+ large_bunch["x"].mean(),
+ rtol=0.1)
+ assert np.isclose(beam_ion.ion_beam["y"].mean(),
+ large_bunch["y"].mean(),
+ rtol=0.1)
# Ion drift tracking properly updates ion positions based on momentum
def test_ion_drift_tracking(self, generate_beam_ion, small_bunch):
@@ -281,29 +338,26 @@ class TestBeamIonElement:
beam_ion.generate_new_ions(small_bunch)
beam_ion.ion_beam["xp"] = np.ones_like(beam_ion.ion_beam["x"])
beam_ion.ion_beam["yp"] = np.ones_like(beam_ion.ion_beam["y"])
-
+
drift_length = 2.0
initial_x = beam_ion.ion_beam["x"].copy()
initial_y = beam_ion.ion_beam["y"].copy()
-
+
beam_ion.track_ions_in_a_drift(drift_length)
-
+
assert np.allclose(beam_ion.ion_beam["x"], initial_x + drift_length)
assert np.allclose(beam_ion.ion_beam["y"], initial_y + drift_length)
# Monitor records ion beam data at specified intervals when enabled
- def test_monitor_recording(self,
- generate_beam_ion,
- small_bunch,
- generate_ion_monitor,
- tmp_path):
- file_name=tmp_path / "test_monitor.hdf5"
+ def test_monitor_recording(self, generate_beam_ion, small_bunch,
+ generate_ion_monitor, tmp_path):
+ file_name = tmp_path / "test_monitor.hdf5"
monitor = generate_ion_monitor(file_name=file_name)
beam_ion = generate_beam_ion()
beam_ion.monitors.append(monitor)
-
+
beam_ion.track(small_bunch)
-
+
assert os.path.exists(file_name)
with hp.File(file_name, 'r') as f:
cond = False
@@ -316,7 +370,7 @@ class TestBeamIonElement:
def test_empty_bunch_handling(self, generate_beam_ion, generate_bunch):
beam_ion = generate_beam_ion()
empty_bunch = generate_bunch(mp_number=0, init_gaussian=False)
-
+
with pytest.raises(ValueError):
beam_ion.generate_new_ions(empty_bunch)
@@ -326,12 +380,13 @@ class TestBeamIonElement:
aperture = IonAperture(X_radius=x_radius, Y_radius=x_radius)
beam_ion = generate_beam_ion()
beam_ion.apertures.append(aperture)
-
+
beam_ion.generate_new_ions(small_bunch)
-
- beam_ion.ion_beam["x"] = np.ones_like(beam_ion.ion_beam["x"]) * (x_radius * 1.1)
+
+ beam_ion.ion_beam["x"] = np.ones_like(
+ beam_ion.ion_beam["x"]) * (x_radius*1.1)
beam_ion.track(beam_ion.ion_beam)
-
+
assert len(beam_ion.ion_beam["x"]) == 0
# Ion clearing removes all particles as expected
@@ -339,27 +394,28 @@ class TestBeamIonElement:
beam_ion = generate_beam_ion()
beam_ion.generate_new_ions(small_bunch)
assert len(beam_ion.ion_beam["x"]) > 0
-
+
beam_ion.clear_ions()
assert len(beam_ion.ion_beam["x"]) == 1
assert beam_ion.ion_beam["x"][0] == 0
-
+
# Tracking with a pre-set cloud
- def test_track_preset_cloud(self, generate_beam_ion, small_bunch, generate_ion_particles):
+ def test_track_preset_cloud(self, generate_beam_ion, small_bunch,
+ generate_ion_particles):
beam_ion = generate_beam_ion()
assert beam_ion.ion_beam.charge == 0
-
+
preset_ion_particles = generate_ion_particles(mp_number=1000)
preset_charge = 1e-9
- preset_ion_particles.generate_as_a_distribution(small_bunch, preset_charge)
+ preset_ion_particles.generate_as_a_distribution(
+ small_bunch, preset_charge)
beam_ion.ion_beam += preset_ion_particles
assert np.isclose(beam_ion.ion_beam.charge, preset_charge)
-
+
beam_ion.generate_ions = False
- assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+ assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp", "yp"])
assert np.isclose(beam_ion.ion_beam.charge, preset_charge)
-
+
beam_ion.generate_ions = True
- assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp","yp"])
+ assert_attr_changed(beam_ion, small_bunch, attrs_changed=["xp", "yp"])
assert beam_ion.ion_beam.charge > preset_charge
-
\ No newline at end of file
diff --git a/tests/unit/tracking/test_element.py b/tests/unit/tracking/test_element.py
index 6084010a573b6e3abd14ad32fba28d96c4f042eb..1048cc0d8fec61c8435812028024a30ad0ad2a12 100644
--- a/tests/unit/tracking/test_element.py
+++ b/tests/unit/tracking/test_element.py
@@ -1,32 +1,44 @@
import numpy as np
import pytest
+
pytestmark = pytest.mark.unit
from scipy.special import factorial
from utility_test_functions import assert_attr_changed
-from mbtrack2 import (Element,
- LongitudinalMap,
- SynchrotronRadiation,
- SkewQuadrupole,
- TransverseMapSector,
- TransverseMap,
- transverse_map_sector_generator)
+
+from mbtrack2 import (
+ Element,
+ LongitudinalMap,
+ SkewQuadrupole,
+ SynchrotronRadiation,
+ TransverseMap,
+ TransverseMapSector,
+ transverse_map_sector_generator,
+)
+
class TestElement:
-
+
def test_parallel_decorator_with_mpi_beam(self, beam_1bunch_mpi):
+
class SubElement(Element):
+
@Element.parallel
def track(self, bunch):
bunch.charge = 1
+
element = SubElement()
element.track(beam_1bunch_mpi)
- assert beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num].charge == pytest.approx(1)
+ assert beam_1bunch_mpi[
+ beam_1bunch_mpi.mpi.bunch_num].charge == pytest.approx(1)
def test_parallel_decorator_with_beam(self, beam_non_uniform):
+
class SubElement(Element):
+
@Element.parallel
def track(self, bunch):
bunch.charge = 1
+
element = SubElement()
element.track(beam_non_uniform)
for i, bunch in enumerate(beam_non_uniform):
@@ -34,16 +46,19 @@ class TestElement:
assert bunch.charge == pytest.approx(1)
else:
assert bunch.charge == pytest.approx(0)
-
+
def test_parallel_decorator_with_bunch(self, small_bunch):
+
class SubElement(Element):
+
@Element.parallel
def track(self, bunch):
bunch.charge = 1
+
element = SubElement()
element.track(small_bunch)
assert small_bunch.charge == pytest.approx(1)
-
+
# Decorator correctly skips track method if bunch is empty
def test_skip_track_if_bunch_empty(self, mocker, generate_bunch):
mock_track = mocker.Mock()
@@ -51,14 +66,14 @@ class TestElement:
empty_bunch = generate_bunch(mp_number=0)
decorated_track(None, empty_bunch)
mock_track.assert_not_called()
-
+
# Decorator calls track method if bunch is not empty
def test_call_track_if_bunch_not_empty(self, mocker, small_bunch):
mock_track = mocker.Mock()
decorated_track = Element.track_bunch_if_non_empty(mock_track)
decorated_track(None, small_bunch)
mock_track.assert_called_once()
-
+
# Decorator respects track_alive flag and calls track method
def test_respect_track_alive_flag(self, mocker, generate_bunch):
mock_track = mocker.Mock()
@@ -70,10 +85,11 @@ class TestElement:
# Track method executes when bunch is not empty and track_alive is True
def test_executes_with_nonempty_bunch(self, small_bunch):
called = []
+
@Element.track_bunch_if_non_empty
def track_method(self, bunch):
called.append(True)
-
+
small_bunch.track_alive = True
track_method(self, small_bunch)
assert called == [True]
@@ -81,10 +97,11 @@ class TestElement:
# Track method executes when track_alive is False regardless of bunch size
def test_executes_when_track_alive_false(self, small_bunch):
called = []
- @Element.track_bunch_if_non_empty
+
+ @Element.track_bunch_if_non_empty
def track_method(self, bunch):
called.append(True)
-
+
small_bunch.track_alive = False
track_method(self, small_bunch)
assert called == [True]
@@ -92,28 +109,33 @@ class TestElement:
# Empty bunch with track_alive=True skips track method execution
def test_skips_empty_bunch(self, generate_bunch):
called = []
+
@Element.track_bunch_if_non_empty
def track_method(self, bunch):
called.append(True)
-
+
empty_bunch = generate_bunch(alive=False)
empty_bunch.track_alive = True
track_method(self, empty_bunch)
assert not called
-
+
+
class TestLongitudinalMap:
# Track a Bunch object using the track method
def test_track_bunch(self, small_bunch, demo_ring):
long_map = LongitudinalMap(demo_ring)
- assert_attr_changed(long_map, small_bunch, attrs_changed=["tau","delta"])
+ assert_attr_changed(long_map,
+ small_bunch,
+ attrs_changed=["tau", "delta"])
+
class TestSynchrotronRadiation:
# SynchrotronRadiation initializes correctly with default switch values
def test_initialization_with_default_switch(self, demo_ring):
sr = SynchrotronRadiation(demo_ring)
- assert np.array_equal(sr.switch, np.ones((3,), dtype=bool))
+ assert np.array_equal(sr.switch, np.ones((3, ), dtype=bool))
# SynchrotronRadiation modifies 'delta', 'xp', and 'yp' attributes of Bunch
def test_modifies_bunch_attributes(self, small_bunch, demo_ring):
@@ -122,22 +144,26 @@ class TestSynchrotronRadiation:
# switch array has all False values, ensuring no changes to Bunch
def test_no_changes_with_all_false_switch(self, small_bunch, demo_ring):
- sr = SynchrotronRadiation(demo_ring, switch=np.zeros((3,), dtype=bool))
+ sr = SynchrotronRadiation(demo_ring,
+ switch=np.zeros((3, ), dtype=bool))
assert_attr_changed(sr, small_bunch, change=False)
-
+
+
class TestSkewQuadrupole:
# Initialize SkewQuadrupole with a positive strength and track a Bunch object
def test_modifies_bunch_attributes(self, small_bunch):
skew_quad = SkewQuadrupole(strength=0.1)
- assert_attr_changed(skew_quad, small_bunch, attrs_changed=["xp","yp"])
-
+ assert_attr_changed(skew_quad, small_bunch, attrs_changed=["xp", "yp"])
+
+
class TestTransverseMapSector:
-
+
@pytest.fixture
def generate_trans_map_sector(demo_ring):
- def generate(phase_diff = np.array([np.pi, np.pi]),
- chro_diff = np.array([0.01, 0.01]),
+
+ def generate(phase_diff=np.array([np.pi, np.pi]),
+ chro_diff=np.array([0.01, 0.01]),
adts=None):
alpha0 = np.array([1.0, 1.0])
beta0 = np.array([1.0, 1.0])
@@ -145,65 +171,85 @@ class TestTransverseMapSector:
alpha1 = np.array([2.0, 2.0])
beta1 = np.array([2.0, 2.0])
dispersion1 = np.array([0.1, 0.1, 0.1, 0.1])
- sector = TransverseMapSector(demo_ring,
- alpha0,
- beta0,
- dispersion0,
- alpha1,
- beta1,
- dispersion1,
- phase_diff,
+ sector = TransverseMapSector(demo_ring,
+ alpha0,
+ beta0,
+ dispersion0,
+ alpha1,
+ beta1,
+ dispersion1,
+ phase_diff,
chro_diff,
adts=adts)
return sector
+
return generate
# Track a Bunch object through TransverseMapSector and ensure coordinates are updated
- def test_track_bunch_coordinates_update(self, generate_trans_map_sector, small_bunch):
+ def test_track_bunch_coordinates_update(self, generate_trans_map_sector,
+ small_bunch):
sector = generate_trans_map_sector()
- assert_attr_changed(sector, small_bunch, attrs_changed=["x", "xp", "y","yp"])
+ assert_attr_changed(sector,
+ small_bunch,
+ attrs_changed=["x", "xp", "y", "yp"])
# Compute chromatic tune advances for a Bunch with non-zero chromaticity differences
- @pytest.mark.parametrize("chro_diff", [(np.array([0.02, 0.03])),
- (np.array([0.02, 0.03, 0.05, 0.06])),
- (np.array([0.02, 0.03, 0.05, 0.06, 0.02, 0.03,])),
- (np.array([0.02, 0.03, 0.05, 0.06, 0.02, 0.03, 0.05, 0.06])),
- (np.array([0.02, 0.03, 0.05, 0.06, 0.02, 0.03, 0.05, 0.06, 0.05, 0.06])),])
- def test_chromatic_tune_advances(self, generate_trans_map_sector, small_bunch, chro_diff):
+ @pytest.mark.parametrize("chro_diff", [
+ (np.array([0.02, 0.03])),
+ (np.array([0.02, 0.03, 0.05, 0.06])),
+ (np.array([
+ 0.02,
+ 0.03,
+ 0.05,
+ 0.06,
+ 0.02,
+ 0.03,
+ ])),
+ (np.array([0.02, 0.03, 0.05, 0.06, 0.02, 0.03, 0.05, 0.06])),
+ (np.array([0.02, 0.03, 0.05, 0.06, 0.02, 0.03, 0.05, 0.06, 0.05, 0.06
+ ])),
+ ])
+ def test_chromatic_tune_advances(self, generate_trans_map_sector,
+ small_bunch, chro_diff):
# chro_diff = np.array([0.02, 0.03])
sector = generate_trans_map_sector(chro_diff=chro_diff)
- tune_advance_x_chro, tune_advance_y_chro = sector._compute_chromatic_tune_advances(small_bunch)
-
+ tune_advance_x_chro, tune_advance_y_chro = sector._compute_chromatic_tune_advances(
+ small_bunch)
+
order = len(chro_diff) // 2
coefs = np.array([1 / factorial(i) for i in range(order + 1)])
coefs[0] = 0
chro_diff = np.concatenate(([0, 0], chro_diff))
- tune_advance_x = np.polynomial.polynomial.Polynomial(chro_diff[::2] * coefs)(small_bunch['delta'])
- tune_advance_y = np.polynomial.polynomial.Polynomial(chro_diff[1::2] * coefs)(small_bunch['delta'])
-
+ tune_advance_x = np.polynomial.polynomial.Polynomial(
+ chro_diff[::2] * coefs)(small_bunch['delta'])
+ tune_advance_y = np.polynomial.polynomial.Polynomial(
+ chro_diff[1::2] * coefs)(small_bunch['delta'])
+
assert np.allclose(tune_advance_x, tune_advance_x_chro)
assert np.allclose(tune_advance_y, tune_advance_y_chro)
-
# Check that adts are taken into account in calculation
- def test_amplitude_dependent_tune_shifts(self, generate_trans_map_sector, small_bunch):
-
+ def test_amplitude_dependent_tune_shifts(self, generate_trans_map_sector,
+ small_bunch):
+
sector_no_adts = generate_trans_map_sector()
- adts=[np.array([1e10, 1e10, 1e10]),
- np.array([1e10, 1e10, 1e10]),
- np.array([1e10, 1e10, 1e10]),
- np.array([1e10, 1e10, 1e10])]
+ adts = [
+ np.array([1e10, 1e10, 1e10]),
+ np.array([1e10, 1e10, 1e10]),
+ np.array([1e10, 1e10, 1e10]),
+ np.array([1e10, 1e10, 1e10])
+ ]
sector_adts = generate_trans_map_sector(adts=adts)
-
- attrs = ["x", "xp", "y","yp"]
+
+ attrs = ["x", "xp", "y", "yp"]
initial_values = {attr: small_bunch[attr].copy() for attr in attrs}
sector_no_adts.track(small_bunch)
no_adts = {attr: small_bunch[attr].copy() for attr in attrs}
-
+
for attr in attrs:
small_bunch[attr] = initial_values[attr]
-
+
sector_adts.track(small_bunch)
adts = {attr: small_bunch[attr].copy() for attr in attrs}
@@ -211,9 +257,10 @@ class TestTransverseMapSector:
assert not np.array_equal(initial_values[attr], no_adts[attr])
assert not np.array_equal(initial_values[attr], adts[attr])
assert not np.array_equal(adts[attr], no_adts[attr])
-
+
+
class TestTransverseMap:
-
+
class Old_TransverseMap(Element):
"""
Transverse map from mbtrack2 0.7.0.
@@ -222,7 +269,7 @@ class TestTransverseMap:
----------
ring : Synchrotron object
"""
-
+
def __init__(self, ring):
self.ring = ring
self.alpha = self.ring.optics.local_alpha
@@ -236,7 +283,7 @@ class TestTransverseMap:
np.poly1d(self.ring.adts[2]),
np.poly1d(self.ring.adts[3]),
]
-
+
@Element.parallel
def track(self, bunch):
"""
@@ -257,7 +304,8 @@ class TestTransverseMap:
phase_advance_y = (
2 * np.pi *
(self.ring.tune[1] + self.ring.chro[1] * bunch["delta"]))
- print("old phase advance ",phase_advance_x[0]/(2*np.pi), phase_advance_y[0]/(2*np.pi))
+ print("old phase advance ", phase_advance_x[0] / (2 * np.pi),
+ phase_advance_y[0] / (2 * np.pi))
else:
Jx = ((self.ring.optics.local_gamma[0] * bunch["x"]**2) +
(2 * self.ring.optics.local_alpha[0] * bunch["x"] *
@@ -275,14 +323,14 @@ class TestTransverseMap:
2 * np.pi *
(self.ring.tune[1] + self.ring.chro[1] * bunch["delta"] +
self.adts_poly[1](Jx) + self.adts_poly[3](Jy)))
-
+
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
matrix = np.zeros((6, 6, len(bunch)), dtype=np.float64)
-
+
# Horizontal
c_x = np.cos(phase_advance_x)
s_x = np.sin(phase_advance_x)
-
+
matrix[0, 0, :] = c_x + self.alpha[0] * s_x
matrix[0, 1, :] = self.beta[0] * s_x
matrix[0, 2, :] = self.dispersion[0]
@@ -292,11 +340,11 @@ class TestTransverseMap:
matrix[2, 2, :] = 1
print(matrix[0, 0, :], matrix[0, 1, :], matrix[0, 2, :])
-
+
# Vertical
c_y = np.cos(phase_advance_y)
s_y = np.sin(phase_advance_y)
-
+
matrix[3, 3, :] = c_y + self.alpha[1] * s_y
matrix[3, 4, :] = self.beta[1] * s_y
matrix[3, 5, :] = self.dispersion[2]
@@ -304,7 +352,7 @@ class TestTransverseMap:
matrix[4, 4, :] = c_y - self.alpha[1] * s_y
matrix[4, 5, :] = self.dispersion[3]
matrix[5, 5, :] = 1
-
+
x = (matrix[0, 0] * bunch["x"] + matrix[0, 1] * bunch["xp"] +
matrix[0, 2] * bunch["delta"])
xp = (matrix[1, 0] * bunch["x"] + matrix[1, 1] * bunch["xp"] +
@@ -313,7 +361,7 @@ class TestTransverseMap:
matrix[3, 5] * bunch["delta"])
yp = (matrix[4, 3] * bunch["y"] + matrix[4, 4] * bunch["yp"] +
matrix[4, 5] * bunch["delta"])
-
+
bunch["x"] = x
bunch["xp"] = xp
bunch["y"] = y
@@ -322,16 +370,16 @@ class TestTransverseMap:
def test_trans_map_base(self, demo_ring, small_bunch):
old_map = self.Old_TransverseMap(demo_ring)
current_map = TransverseMap(demo_ring)
-
- attrs = ["x", "xp", "y","yp"]
+
+ attrs = ["x", "xp", "y", "yp"]
initial_values = {attr: small_bunch[attr].copy() for attr in attrs}
old_map.track(small_bunch)
old = {attr: small_bunch[attr].copy() for attr in attrs}
-
+
for attr in attrs:
small_bunch[attr] = initial_values[attr]
-
+
current_map.track(small_bunch)
current = {attr: small_bunch[attr].copy() for attr in attrs}
@@ -339,12 +387,12 @@ class TestTransverseMap:
assert not np.array_equal(initial_values[attr], current[attr])
assert not np.array_equal(initial_values[attr], old[attr])
assert np.allclose(current[attr], old[attr])
-
+
def test_trans_map_adts(self, ring_with_at_lattice, small_bunch):
ring_with_at_lattice.get_adts()
# Needed as old map was not handling dispersion properly!
- ring_with_at_lattice.optics.local_dispersion = np.array([0,0,0,0])
+ ring_with_at_lattice.optics.local_dispersion = np.array([0, 0, 0, 0])
old_map = self.Old_TransverseMap(ring_with_at_lattice)
current_map = TransverseMap(ring_with_at_lattice)
@@ -358,15 +406,15 @@ class TestTransverseMap:
assert np.array_equal(old_map.dispersion, current_map.dispersion0)
assert np.array_equal(old_map.dispersion, current_map.dispersion1)
- attrs = ["x", "xp", "y","yp"]
+ attrs = ["x", "xp", "y", "yp"]
initial_values = {attr: small_bunch[attr].copy() for attr in attrs}
old_map.track(small_bunch)
old = {attr: small_bunch[attr].copy() for attr in attrs}
-
+
for attr in attrs:
small_bunch[attr] = initial_values[attr]
-
+
current_map.track(small_bunch)
current = {attr: small_bunch[attr].copy() for attr in attrs}
@@ -375,45 +423,54 @@ class TestTransverseMap:
assert not np.array_equal(initial_values[attr], old[attr])
np.testing.assert_allclose(current[attr], old[attr])
+
class TestTransverseMapSectorGenerator:
# Generate sectors with local optics values when use_local_values is True
def test_local_optics_values(self, demo_ring):
positions = np.array([0, 25, 50, 75])
sectors = transverse_map_sector_generator(demo_ring, positions)
-
+
assert len(sectors) == len(positions)
for sector in sectors:
assert np.array_equal(sector.alpha0, demo_ring.optics.local_alpha)
assert np.array_equal(sector.beta0, demo_ring.optics.local_beta)
- assert np.array_equal(sector.dispersion0, demo_ring.optics.local_dispersion)
+ assert np.array_equal(sector.dispersion0,
+ demo_ring.optics.local_dispersion)
# Generate sectors with AT lattice optics when use_local_values is False
def test_at_lattice_optics(self, ring_with_at_lattice):
positions = np.array([0, 25, 50, 75])
- sectors = transverse_map_sector_generator(ring_with_at_lattice, positions)
-
+ sectors = transverse_map_sector_generator(ring_with_at_lattice,
+ positions)
+
assert len(sectors) == len(positions)
for i, sector in enumerate(sectors):
- assert np.array_equal(sector.alpha0, ring_with_at_lattice.optics.alpha(positions[i]))
- assert np.array_equal(sector.beta0, ring_with_at_lattice.optics.beta(positions[i]))
+ assert np.array_equal(
+ sector.alpha0, ring_with_at_lattice.optics.alpha(positions[i]))
+ assert np.array_equal(
+ sector.beta0, ring_with_at_lattice.optics.beta(positions[i]))
# Calculate phase differences between consecutive positions
def test_phase_differences(self, ring_with_at_lattice):
positions = np.array([0, 25, 50])
ring_with_at_lattice.optics.use_local_values = False
- sectors = transverse_map_sector_generator(ring_with_at_lattice, positions)
-
+ sectors = transverse_map_sector_generator(ring_with_at_lattice,
+ positions)
+
for i, sector in enumerate(sectors):
- if i < len(positions)-1:
- expected_phase = ring_with_at_lattice.optics.mu(positions[i+1]) - ring_with_at_lattice.optics.mu(positions[i])
- assert np.allclose(sector.tune_diff * (2 * np.pi), expected_phase)
+ if i < len(positions) - 1:
+ expected_phase = ring_with_at_lattice.optics.mu(
+ positions[i + 1]) - ring_with_at_lattice.optics.mu(
+ positions[i])
+ assert np.allclose(sector.tune_diff * (2 * np.pi),
+ expected_phase)
# Compute chromaticity differences between sectors
def test_chromaticity_differences(self, demo_ring):
positions = np.array([0, 50])
sectors = transverse_map_sector_generator(demo_ring, positions)
-
+
expected_chro = np.asarray(demo_ring.chro) / len(positions)
assert np.allclose(sectors[0].chro_diff, expected_chro)
@@ -422,6 +479,6 @@ class TestTransverseMapSectorGenerator:
demo_ring.adts = np.array([1.0, 1.0, 1.0, 1.0])
positions = np.array([0, 50])
sectors = transverse_map_sector_generator(demo_ring, positions)
-
+
expected_adts = demo_ring.adts / len(positions)
- assert np.allclose(sectors[0].adts_poly[0].coef, expected_adts[0])
\ No newline at end of file
+ assert np.allclose(sectors[0].adts_poly[0].coef, expected_adts[0])
diff --git a/tests/unit/tracking/test_emfields.py b/tests/unit/tracking/test_emfields.py
index 40cbea56b12e8a09e47af0238d7d61c2926974f1..ff95fb06605d675d11618cd341df8b4699cff93d 100644
--- a/tests/unit/tracking/test_emfields.py
+++ b/tests/unit/tracking/test_emfields.py
@@ -1,18 +1,19 @@
import numpy as np
+import pytest
from mbtrack2.tracking.emfields import (
- _wofz,
- _sqrt_sig,
- _efieldn_mit,
_efieldn_linearized,
- efieldn_gauss_round,
+ _efieldn_mit,
+ _sqrt_sig,
+ _wofz,
add_sigma_check,
+ efieldn_gauss_round,
get_displaced_efield,
)
-import pytest
pytestmark = pytest.mark.unit
+
class Test_particles_electromagnetic_fields:
def test_wofz_return_float(self):
@@ -24,9 +25,8 @@ class Test_particles_electromagnetic_fields:
val = _sqrt_sig(1.0, 1.0)
assert isinstance(val, float)
assert val >= 0
-
- @pytest.mark.parametrize("func",[(_efieldn_mit),
- (efieldn_gauss_round),
+
+ @pytest.mark.parametrize("func", [(_efieldn_mit), (efieldn_gauss_round),
(_efieldn_linearized)])
def test_efieldn_return_float(self, func):
ex, ey = func(1.0, 1.0, 1.0, 1.0)
@@ -34,7 +34,9 @@ class Test_particles_electromagnetic_fields:
assert isinstance(ey, float)
# Maintains original function behavior when sig_x is greater than sig_y
- def test_add_sigma_check_maintain_original_behavior_when_sig_x_greater_than_sig_y(self):
+ def test_add_sigma_check_maintain_original_behavior_when_sig_x_greater_than_sig_y(
+ self):
+
def mock_efieldn(x, y, sig_x, sig_y):
return x + sig_x, y + sig_y
@@ -44,6 +46,7 @@ class Test_particles_electromagnetic_fields:
# Exchanges x and y when sig_x is less than sig_y
def test_add_sigma_check_exchange_x_y_when_sig_x_less_than_sig_y(self):
+
def mock_efieldn(x, y, sig_x, sig_y):
return x + sig_x, y + sig_y
@@ -54,50 +57,56 @@ class Test_particles_electromagnetic_fields:
# Applies round beam field formula when sig_x is close to sig_y
def test_add_sigma_check_apply_round_beam_formula_when_sigmas_close(self):
wrapped_function = add_sigma_check(efieldn_gauss_round)
- result = wrapped_function(np.array([1.0]), np.array([1.0]), 1.0, 1.00001)
- assert np.allclose(result, efieldn_gauss_round(np.array([1.0]), np.array([1.0]), 1.0, 1.00001))
+ result = wrapped_function(np.array([1.0]), np.array([1.0]), 1.0,
+ 1.00001)
+ assert np.allclose(
+ result,
+ efieldn_gauss_round(np.array([1.0]), np.array([1.0]), 1.0,
+ 1.00001))
# Returns zero fields when sig_x and sig_y are both close to zero
def test_add_sigma_check_zero_fields_when_sigmas_close_to_zero(self):
wrapped_function = add_sigma_check(efieldn_gauss_round)
- result = wrapped_function(np.array([1.0]), np.array([1.0]), 1e-11, 1e-11)
+ result = wrapped_function(np.array([1.0]), np.array([1.0]), 1e-11,
+ 1e-11)
assert np.allclose(result, (np.zeros(1), np.zeros(1)))
-
+
def test_add_sigma_check_empty_arrays(self):
# Define a mock efieldn function
def mock_efieldn(x, y, sig_x, sig_y):
return np.zeros_like(x), np.zeros_like(y)
-
+
# Wrap the mock function with add_sigma_check
wrapped_function = add_sigma_check(mock_efieldn)
-
+
# Create empty arrays for x and y
x = np.array([])
y = np.array([])
sig_x = 1.0
sig_y = 1.0
-
+
# Call the wrapped function
en_x, en_y = wrapped_function(x, y, sig_x, sig_y)
-
+
# Assert that the output is also empty arrays
assert en_x.size == 0
assert en_y.size == 0
-
+
# Computes electric field for typical Gaussian charge distribution
def test_get_displaced_efield_return_shape(self):
+
def mock_efieldn(x, y, sig_x, sig_y):
return np.ones_like(x), np.ones_like(y)
-
+
xr = np.array([1.0, 2.0, 3.0])
yr = np.array([1.0, 2.0, 3.0])
sig_x = 2.0
sig_y = 1.0
mean_x = 0.0
mean_y = 0.0
-
- en_x, en_y = get_displaced_efield(mock_efieldn, xr, yr, sig_x, sig_y, mean_x, mean_y)
-
+
+ en_x, en_y = get_displaced_efield(mock_efieldn, xr, yr, sig_x, sig_y,
+ mean_x, mean_y)
+
assert np.allclose(en_x, [1.0, 1.0, 1.0])
assert np.allclose(en_y, [1.0, 1.0, 1.0])
-
diff --git a/tests/unit/tracking/test_excite.py b/tests/unit/tracking/test_excite.py
index f698a18ccbf1c41a59d19d6f1c26ad09b730cde8..c7721807a5be65e7b4d27eb66c02a0dbf9ab303e 100644
--- a/tests/unit/tracking/test_excite.py
+++ b/tests/unit/tracking/test_excite.py
@@ -1,11 +1,15 @@
import numpy as np
import pytest
+
pytestmark = pytest.mark.unit
-from mbtrack2 import Sweep
from utility_test_functions import assert_attr_changed
+from mbtrack2 import Sweep
+
+
@pytest.fixture
def generate_sweep(demo_ring):
+
def generate(ring=demo_ring,
f0=1e3,
f1=2e3,
@@ -21,34 +25,44 @@ def generate_sweep(demo_ring):
plane=plane,
bunch_to_sweep=bunch_to_sweep)
return sweep
+
return generate
+
class TestSweep:
# Sweep applies frequency chirp correctly to a single bunch
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("tau","delta")])
- def test_single_bunch_chirp(self, generate_sweep, small_bunch, plane, attr):
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("tau", "delta")])
+ def test_single_bunch_chirp(self, generate_sweep, small_bunch, plane,
+ attr):
sweep = generate_sweep(plane=plane)
assert_attr_changed(sweep, small_bunch, attrs_changed=[attr])
# Sweep applies frequency chirp correctly to all bunches in a beam
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("tau","delta")])
- def test_all_bunches_chirp(self, generate_sweep, beam_uniform, plane, attr):
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("tau", "delta")])
+ def test_all_bunches_chirp(self, generate_sweep, beam_uniform, plane,
+ attr):
sweep = generate_sweep(plane=plane)
assert_attr_changed(sweep, beam_uniform, attrs_changed=[attr])
# Sweep applies frequency chirp correctly to all bunches in a beam
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("tau","delta")])
- def test_beam_chirp_single_bunch(self, generate_sweep, beam_uniform, plane, attr):
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("tau", "delta")])
+ def test_beam_chirp_single_bunch(self, generate_sweep, beam_uniform, plane,
+ attr):
sweep = generate_sweep(plane=plane, bunch_to_sweep=5)
initial_attr = [bunch[attr].copy() for bunch in beam_uniform]
sweep.track(beam_uniform)
for i, bunch in enumerate(beam_uniform):
if i == 5:
- assert not np.allclose(initial_attr[i], bunch[attr]), f"Chirp not applied correctly to bunch {i}"
+ assert not np.allclose(
+ initial_attr[i],
+ bunch[attr]), f"Chirp not applied correctly to bunch {i}"
else:
- assert np.allclose(initial_attr[i], bunch[attr]), f"Chirp applied to bunch {i}"
-
+ assert np.allclose(initial_attr[i],
+ bunch[attr]), f"Chirp applied to bunch {i}"
# Plot method generates a correct sweep voltage plot
def test_plot_sweep_voltage(self, generate_sweep):
@@ -57,17 +71,19 @@ class TestSweep:
assert fig is not None, "Plot method did not return a figure"
# Sweep tracks correctly when MPI is enabled
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("tau","delta")])
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("tau", "delta")])
def test_mpi_tracking(self, generate_sweep, beam_1bunch_mpi, plane, attr):
sweep = generate_sweep(plane=plane)
bunch = beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num]
initial_attr = bunch[attr].copy()
sweep.track(beam_1bunch_mpi)
- assert not np.allclose(initial_attr, bunch[attr]), "Chirp not applied correctly with MPI"
+ assert not np.allclose(
+ initial_attr, bunch[attr]), "Chirp not applied correctly with MPI"
# Sweep correctly resets count after completing a full sweep
def test_count_reset_after_full_sweep(self, generate_sweep, small_bunch):
sweep = generate_sweep()
for _ in range(sweep.N): # Complete one full sweep cycle
sweep.track(small_bunch)
- assert sweep.count == 0, "Count did not reset after completing a full sweep"
\ No newline at end of file
+ assert sweep.count == 0, "Count did not reset after completing a full sweep"
diff --git a/tests/unit/tracking/test_ibs.py b/tests/unit/tracking/test_ibs.py
index c119720053daf3d3344ea70d3a6e34a4cc80fd6f..2265a03f2a81ea02560e3dfa39c278549581ca73 100644
--- a/tests/unit/tracking/test_ibs.py
+++ b/tests/unit/tracking/test_ibs.py
@@ -1,11 +1,15 @@
-from mbtrack2 import IntrabeamScattering
import numpy as np
import pytest
+
+from mbtrack2 import IntrabeamScattering
+
pytestmark = pytest.mark.unit
from utility_test_functions import assert_attr_changed
+
@pytest.fixture
def generate_ibs(ring_with_at_lattice):
+
def generate(ring=ring_with_at_lattice,
model="CIMP",
n_points=100,
@@ -15,16 +19,20 @@ def generate_ibs(ring_with_at_lattice):
n_points=n_points,
n_bin=n_bin)
return ibs
+
return generate
+
class TestIntrabeamScattering:
-
+
# Track a bunch and validate the momentum changes
- @pytest.mark.parametrize("model",[("Bane"),("CIMP"), ("PM"), ("PS")])
+ @pytest.mark.parametrize("model", [("Bane"), ("CIMP"), ("PM"), ("PS")])
def test_track_model_change_attr(self, small_bunch, generate_ibs, model):
ibs = generate_ibs(model=model)
if model == "Bane":
- assert_attr_changed(ibs, small_bunch, attrs_changed=["xp", "delta"])
+ assert_attr_changed(ibs,
+ small_bunch,
+ attrs_changed=["xp", "delta"])
else:
assert_attr_changed(ibs, small_bunch)
@@ -44,13 +52,13 @@ class TestIntrabeamScattering:
small_bunch.alive[:] = False
ibs = generate_ibs()
assert_attr_changed(ibs, small_bunch, change=False)
-
+
# Handle a case where the bunch has lost some macro-particles
def test_partial_empty_bunch_handling(self, small_bunch, generate_ibs):
small_bunch.alive[0:5] = False
ibs = generate_ibs()
assert_attr_changed(ibs, small_bunch)
-
+
# Test the warning when lattice file is not loaded and optics are approximated
def test_warning_for_approximated_optics(self, demo_ring):
demo_ring.optics.use_local_values = True
@@ -58,17 +66,20 @@ class TestIntrabeamScattering:
IntrabeamScattering(ring=demo_ring, model="CIMP")
# Validate the computation of growth rates for each model
- def test_growth_rate_computation(self, demo_ring, small_bunch, generate_ibs):
+ def test_growth_rate_computation(self, demo_ring, small_bunch,
+ generate_ibs):
for model in ["PS", "PM", "Bane", "CIMP"]:
ibs = generate_ibs(model=model)
ibs.initialize(small_bunch)
if model in ["PS", "PM"]:
vabq, v1aq, v1bq = ibs.scatter()
- T_x, T_y, T_p = ibs.get_scatter_T(vabq=vabq, v1aq=v1aq, v1bq=v1bq)
+ T_x, T_y, T_p = ibs.get_scatter_T(vabq=vabq,
+ v1aq=v1aq,
+ v1bq=v1bq)
elif model == "Bane":
gval = ibs.scatter()
T_x, T_y, T_p = ibs.get_scatter_T(gval=gval)
elif model == "CIMP":
g_ab, g_ba = ibs.scatter()
T_x, T_y, T_p = ibs.get_scatter_T(g_ab=g_ab, g_ba=g_ba)
- assert T_x >= 0 and T_y >= 0 and T_p >= 0
\ No newline at end of file
+ assert T_x >= 0 and T_y >= 0 and T_p >= 0
diff --git a/tests/unit/tracking/test_parallel.py b/tests/unit/tracking/test_parallel.py
index dd75e633103bec871170206a77c91fafd7785f6e..76e0eb1721f32227679660ddb2d05435b74edeb7 100644
--- a/tests/unit/tracking/test_parallel.py
+++ b/tests/unit/tracking/test_parallel.py
@@ -1,34 +1,41 @@
import pytest
+
pytestmark = pytest.mark.unit
import numpy as np
-from mbtrack2.tracking.parallel import Mpi
from mpi4py import MPI
+from mbtrack2.tracking.parallel import Mpi
+
+
@pytest.fixture
def mpi_size1():
filling_pattern = np.array([True])
return Mpi(filling_pattern)
+
@pytest.fixture
def generate_mpi_mock(mocker):
- def generate(filling_pattern,
- rank=0):
-
+
+ def generate(filling_pattern, rank=0):
+
size = filling_pattern.sum()
mocker.patch.object(MPI, 'COMM_WORLD')
MPI.COMM_WORLD.Get_size.return_value = size
MPI.COMM_WORLD.Get_rank.return_value = rank
return Mpi(filling_pattern)
+
return generate
+
@pytest.fixture
def mpi_mock_size4_fill6(generate_mpi_mock):
- filling_pattern = np.ones((6,), dtype=bool)
+ filling_pattern = np.ones((6, ), dtype=bool)
filling_pattern[1] = False
filling_pattern[4] = False
return generate_mpi_mock(filling_pattern)
+
class TestMpi:
# Verify table creation
@@ -39,7 +46,8 @@ class TestMpi:
assert mpi_size1.previous_bunch == 0
# Verify correct rank-to-bunch and bunch-to-rank mappings
- def test_rank_to_bunch_and_bunch_to_rank_mappings(self, mpi_mock_size4_fill6):
+ def test_rank_to_bunch_and_bunch_to_rank_mappings(self,
+ mpi_mock_size4_fill6):
assert mpi_mock_size4_fill6.rank_to_bunch(0) == 0
assert mpi_mock_size4_fill6.bunch_to_rank(0) == 0
assert mpi_mock_size4_fill6.rank_to_bunch(1) == 2
@@ -55,45 +63,68 @@ class TestMpi:
assert mpi_mock_size4_fill6.rank == 0
assert mpi_mock_size4_fill6.next_bunch == 1
assert mpi_mock_size4_fill6.previous_bunch == 3
-
- @pytest.mark.parametrize("dim", [("tau"), ("delta"), ("x"), ("xp"), ("y"), ("yp")])
+
+ @pytest.mark.parametrize("dim", [("tau"), ("delta"), ("x"), ("xp"), ("y"),
+ ("yp")])
def test_share_distributions(self, beam_1bunch_mpi, dim):
n_bin = 75
- beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi, dimensions=dim, n_bin=n_bin)
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_center").shape == (beam_1bunch_mpi.mpi.size, n_bin)
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_profile").shape == (beam_1bunch_mpi.mpi.size, n_bin)
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_sorted_index").shape == (len(beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num]),)
+ beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi,
+ dimensions=dim,
+ n_bin=n_bin)
+ assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_center").shape == (
+ beam_1bunch_mpi.mpi.size, n_bin)
+ assert beam_1bunch_mpi.mpi.__getattribute__(
+ dim + "_profile").shape == (beam_1bunch_mpi.mpi.size, n_bin)
+ assert beam_1bunch_mpi.mpi.__getattribute__(
+ dim + "_sorted_index").shape == (len(
+ beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num]), )
assert beam_1bunch_mpi.mpi.charge_per_mp_all is not None
def test_share_distributions_multidim(self, beam_1bunch_mpi):
- n_bin = [75,60,22]
- dims = ["x","yp","tau"]
- beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi, dimensions=dims, n_bin=n_bin)
+ n_bin = [75, 60, 22]
+ dims = ["x", "yp", "tau"]
+ beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi,
+ dimensions=dims,
+ n_bin=n_bin)
assert beam_1bunch_mpi.mpi.charge_per_mp_all is not None
for i, dim in enumerate(dims):
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_center").shape == (beam_1bunch_mpi.mpi.size, n_bin[i])
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_profile").shape == (beam_1bunch_mpi.mpi.size, n_bin[i])
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_sorted_index").shape == (len(beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num]),)
+ assert beam_1bunch_mpi.mpi.__getattribute__(
+ dim + "_center").shape == (beam_1bunch_mpi.mpi.size, n_bin[i])
+ assert beam_1bunch_mpi.mpi.__getattribute__(
+ dim + "_profile").shape == (beam_1bunch_mpi.mpi.size, n_bin[i])
+ assert beam_1bunch_mpi.mpi.__getattribute__(
+ dim + "_sorted_index").shape == (len(
+ beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num]), )
def test_share_distributions_nobunch(self, beam_1bunch_mpi):
n_bin = 75
dim = "x"
beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num].alive[:] = False
- beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi, dimensions=dim, n_bin=n_bin)
+ beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi,
+ dimensions=dim,
+ n_bin=n_bin)
assert beam_1bunch_mpi.mpi.charge_per_mp_all is not None
- assert np.allclose(beam_1bunch_mpi.mpi.__getattribute__(dim + "_center"), np.zeros((n_bin, ), dtype=np.int64))
- assert np.allclose(beam_1bunch_mpi.mpi.__getattribute__(dim + "_profile"), np.zeros((n_bin, ), dtype=np.float64))
- assert beam_1bunch_mpi.mpi.__getattribute__(dim + "_sorted_index") is None
-
+ assert np.allclose(
+ beam_1bunch_mpi.mpi.__getattribute__(dim + "_center"),
+ np.zeros((n_bin, ), dtype=np.int64))
+ assert np.allclose(
+ beam_1bunch_mpi.mpi.__getattribute__(dim + "_profile"),
+ np.zeros((n_bin, ), dtype=np.float64))
+ assert beam_1bunch_mpi.mpi.__getattribute__(dim +
+ "_sorted_index") is None
def test_share_means(self, beam_1bunch_mpi):
beam_1bunch_mpi.mpi.share_means(beam_1bunch_mpi)
assert beam_1bunch_mpi.mpi.charge_all is not None
- assert beam_1bunch_mpi.mpi.mean_all.shape == (beam_1bunch_mpi.mpi.size, 6)
- assert beam_1bunch_mpi.mpi.mean_all[0,:] == pytest.approx(beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num].mean)
+ assert beam_1bunch_mpi.mpi.mean_all.shape == (beam_1bunch_mpi.mpi.size,
+ 6)
+ assert beam_1bunch_mpi.mpi.mean_all[0, :] == pytest.approx(
+ beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num].mean)
def test_share_stds(self, beam_1bunch_mpi):
beam_1bunch_mpi.mpi.share_stds(beam_1bunch_mpi)
assert beam_1bunch_mpi.mpi.charge_all is not None
- assert beam_1bunch_mpi.mpi.std_all.shape == (beam_1bunch_mpi.mpi.size, 6)
- assert beam_1bunch_mpi.mpi.std_all[0,:] == pytest.approx(beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num].std)
\ No newline at end of file
+ assert beam_1bunch_mpi.mpi.std_all.shape == (beam_1bunch_mpi.mpi.size,
+ 6)
+ assert beam_1bunch_mpi.mpi.std_all[0, :] == pytest.approx(
+ beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num].std)
diff --git a/tests/unit/tracking/test_particle.py b/tests/unit/tracking/test_particle.py
index 2f0a174e8bed0ef2031270c3194cb38ced935567..836640b2914e7f7733c73e471a7d186b1b2a455d 100644
--- a/tests/unit/tracking/test_particle.py
+++ b/tests/unit/tracking/test_particle.py
@@ -1,9 +1,12 @@
-import numpy as np
import matplotlib.pyplot as plt
+import numpy as np
import pytest
+
pytestmark = pytest.mark.unit
-from scipy.constants import e, m_p, c
-from mbtrack2 import Particle, Bunch, Beam
+from scipy.constants import c, e, m_p
+
+from mbtrack2 import Beam, Bunch, Particle
+
class TestParticle:
@@ -12,38 +15,45 @@ class TestParticle:
particle = Particle(mass=m_p, charge=e)
expected_E_rest = m_p * c**2 / e
assert particle.E_rest == expected_E_rest
-
+
+
@pytest.fixture
def generate_bunch(demo_ring):
- def generate(ring=demo_ring,
- mp_number=1e3,
- current=1e-3,
- track_alive=True,
- alive=True,
- load_from_file=None,
- load_suffix=None,
- init_gaussian=True,
- ):
+
+ def generate(
+ ring=demo_ring,
+ mp_number=1e3,
+ current=1e-3,
+ track_alive=True,
+ alive=True,
+ load_from_file=None,
+ load_suffix=None,
+ init_gaussian=True,
+ ):
bunch = Bunch(ring=ring,
- mp_number=mp_number,
- current=current,
- track_alive=track_alive,
- alive=alive,
- load_from_file=load_from_file,
- load_suffix=load_suffix)
+ mp_number=mp_number,
+ current=current,
+ track_alive=track_alive,
+ alive=alive,
+ load_from_file=load_from_file,
+ load_suffix=load_suffix)
if init_gaussian:
bunch.init_gaussian()
return bunch
+
return generate
+
@pytest.fixture
def small_bunch(generate_bunch):
return generate_bunch(mp_number=10)
+
@pytest.fixture
def large_bunch(generate_bunch):
return generate_bunch(mp_number=1e5)
+
class TestBunch:
# Calculate and verify the mean, std, skew, and kurtosis of particle positions
@@ -67,7 +77,7 @@ class TestBunch:
bunch = generate_bunch(mp_number=0)
assert bunch.mp_number == 0
assert bunch.is_empty == True
-
+
def test_drop_zero_macro_particles(self, generate_bunch):
bunch = generate_bunch(mp_number=1)
bunch.alive[0] = False
@@ -77,24 +87,30 @@ class TestBunch:
# Verify the behavior of init_gaussian with custom covariance and mean
def test_init_gaussian_custom_cov_mean(self, large_bunch):
-
+
custom_cov = np.eye(6) * 0.5
custom_mean = np.ones(6) * 2.0
-
+
large_bunch.init_gaussian(cov=custom_cov, mean=custom_mean)
-
+
assert np.allclose(large_bunch.mean, custom_mean, atol=1e-1)
- assert np.allclose(np.cov(large_bunch.particles['x'], large_bunch.particles['xp']), custom_cov[:2, :2], atol=1e-1)
+ assert np.allclose(np.cov(large_bunch.particles['x'],
+ large_bunch.particles['xp']),
+ custom_cov[:2, :2],
+ atol=1e-1)
def test_bunch_values(self, small_bunch, demo_ring):
mp_number = small_bunch.mp_number
current = small_bunch.current
assert len(small_bunch) == mp_number
- np.testing.assert_allclose(small_bunch.alive, np.ones((mp_number,), dtype=bool))
+ np.testing.assert_allclose(small_bunch.alive,
+ np.ones((mp_number, ), dtype=bool))
assert pytest.approx(small_bunch.charge) == current * demo_ring.T0
- assert pytest.approx(small_bunch.charge_per_mp) == current * demo_ring.T0 / mp_number
- assert pytest.approx(small_bunch.particle_number) == current * demo_ring.T0 / e
+ assert pytest.approx(
+ small_bunch.charge_per_mp) == current * demo_ring.T0 / mp_number
+ assert pytest.approx(
+ small_bunch.particle_number) == current * demo_ring.T0 / e
assert small_bunch.is_empty == False
def test_bunch_magic(self, generate_bunch):
@@ -108,41 +124,47 @@ class TestBunch:
charge_init = small_bunch.charge
small_bunch.alive[0] = False
assert len(small_bunch) == small_bunch.mp_number - 1
- assert pytest.approx(small_bunch.charge) == charge_init * len(small_bunch) / small_bunch.mp_number
+ assert pytest.approx(
+ small_bunch.charge
+ ) == charge_init * len(small_bunch) / small_bunch.mp_number
def test_bunch_init_gauss(self, large_bunch):
- large_bunch.init_gaussian(mean=np.ones((6,)))
- np.testing.assert_allclose(large_bunch.mean, np.ones((6,)), rtol=1e-2)
+ 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(self, small_bunch, generate_bunch, tmp_path):
small_bunch["x"] += 1
small_bunch.save(str(tmp_path / "test"))
-
+
mybunch2 = generate_bunch(mp_number=1, current=1e-5)
mybunch2.load(str(tmp_path / "test.hdf5"))
-
+
assert small_bunch.mp_number == mybunch2.mp_number
assert pytest.approx(small_bunch.charge) == mybunch2.charge
for label in small_bunch:
np.testing.assert_allclose(small_bunch[label], mybunch2[label])
def test_bunch_stats(self, demo_ring, large_bunch):
- 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_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)
-
- @pytest.mark.parametrize('n_bin',
- [(75),
- (1),
- (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)
+
+ @pytest.mark.parametrize('n_bin', [(75), (1), (2)])
def test_bunch_binning(self, small_bunch, n_bin):
- (bins, sorted_index, profile, center) = small_bunch.binning(n_bin=n_bin)
- profile0 = np.zeros((len(bins)-1,))
+ (bins, sorted_index, profile,
+ center) = small_bunch.binning(n_bin=n_bin)
+ profile0 = np.zeros((len(bins) - 1, ))
for i, val in enumerate(sorted_index):
- assert bins[val] <= small_bunch["tau"][i] <= bins[val+1]
+ assert bins[val] <= small_bunch["tau"][i] <= bins[val + 1]
profile0[val] += 1
np.testing.assert_allclose(profile0, profile)
@@ -154,150 +176,220 @@ class TestBunch:
def test_bunch_emittance(self, generate_bunch, demo_ring):
mp_number = 1_000_000
mybunch = generate_bunch(mp_number=mp_number, track_alive=False)
-
- np.testing.assert_allclose(mybunch.emit[0], demo_ring.emit[0], rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated')
- np.testing.assert_allclose(mybunch.emit[1], demo_ring.emit[1], rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated')
-
- np.testing.assert_allclose(mybunch.emit[0], mybunch.cs_invariant[0]/2, rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only')
- np.testing.assert_allclose(mybunch.emit[1], mybunch.cs_invariant[1]/2, rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only')
-
+
+ np.testing.assert_allclose(
+ mybunch.emit[0],
+ demo_ring.emit[0],
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated'
+ )
+ np.testing.assert_allclose(
+ mybunch.emit[1],
+ demo_ring.emit[1],
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated'
+ )
+
+ np.testing.assert_allclose(
+ mybunch.emit[0],
+ mybunch.cs_invariant[0] / 2,
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only'
+ )
+ np.testing.assert_allclose(
+ mybunch.emit[1],
+ mybunch.cs_invariant[1] / 2,
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only'
+ )
+
def test_bunch_emittance_with_dispersion(self, generate_bunch, demo_ring):
mp_number = 1_000_000
mybunch = generate_bunch(mp_number=mp_number, track_alive=False)
- np.testing.assert_allclose(mybunch.emit[0], demo_ring.emit[0], rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated')
- np.testing.assert_allclose(mybunch.emit[1], demo_ring.emit[1], rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated')
-
- np.testing.assert_allclose(mybunch.emit[0], mybunch.cs_invariant[0]/2, rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only')
- np.testing.assert_allclose(mybunch.emit[1], mybunch.cs_invariant[1]/2, rtol=1e-2, atol=0,
- err_msg=f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only')
+ np.testing.assert_allclose(
+ mybunch.emit[0],
+ demo_ring.emit[0],
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated'
+ )
+ np.testing.assert_allclose(
+ mybunch.emit[1],
+ demo_ring.emit[1],
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated'
+ )
+
+ np.testing.assert_allclose(
+ mybunch.emit[0],
+ mybunch.cs_invariant[0] / 2,
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only'
+ )
+ np.testing.assert_allclose(
+ mybunch.emit[1],
+ mybunch.cs_invariant[1] / 2,
+ rtol=1e-2,
+ atol=0,
+ err_msg=
+ f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only'
+ )
@pytest.fixture
def generate_beam(demo_ring, generate_bunch):
+
def generate(ring=demo_ring,
- filling_pattern=None,
- current_per_bunch=1e-3,
+ filling_pattern=None,
+ current_per_bunch=1e-3,
mp_per_bunch=10,
track_alive=True,
mpi=False):
beam = Beam(ring=ring)
if filling_pattern is None:
- filling_pattern = np.ones((ring.h,), dtype=bool)
- beam.init_beam(filling_pattern=filling_pattern,
- current_per_bunch=current_per_bunch,
+ filling_pattern = np.ones((ring.h, ), dtype=bool)
+ beam.init_beam(filling_pattern=filling_pattern,
+ current_per_bunch=current_per_bunch,
mp_per_bunch=mp_per_bunch,
track_alive=track_alive,
mpi=mpi)
return beam
+
return generate
+
@pytest.fixture
def beam_uniform(generate_beam):
return generate_beam()
+
@pytest.fixture
def beam_non_uniform(generate_beam, demo_ring):
- filling_pattern = np.ones((demo_ring.h,), dtype=bool)
+ filling_pattern = np.ones((demo_ring.h, ), dtype=bool)
filling_pattern[4] = False
filling_pattern[-3:] = False
return generate_beam(filling_pattern=filling_pattern)
+
@pytest.fixture
def beam_1bunch_mpi(generate_beam, demo_ring):
- filling_pattern = np.zeros((demo_ring.h,), dtype=bool)
+ filling_pattern = np.zeros((demo_ring.h, ), dtype=bool)
filling_pattern[0] = True
return generate_beam(filling_pattern=filling_pattern, mpi=True)
+
class TestBeam:
- @pytest.mark.parametrize("n_bunch", [(1),(5),(10),(20)])
+ @pytest.mark.parametrize("n_bunch", [(1), (5), (10), (20)])
def test_initialize_beam(self, generate_beam, demo_ring, n_bunch):
- filling_pattern = np.zeros((demo_ring.h,), dtype=bool)
+ filling_pattern = np.zeros((demo_ring.h, ), dtype=bool)
filling_pattern[0:n_bunch] = True
beam = generate_beam(filling_pattern=filling_pattern)
assert len(beam) == n_bunch
- @pytest.mark.parametrize("n_bunch", [(1),(5),(10),(20)])
- def test_calculate_total_beam_properties(self, generate_beam, demo_ring, n_bunch):
- filling_pattern = np.zeros((demo_ring.h,), dtype=bool)
+ @pytest.mark.parametrize("n_bunch", [(1), (5), (10), (20)])
+ def test_calculate_total_beam_properties(self, generate_beam, demo_ring,
+ n_bunch):
+ filling_pattern = np.zeros((demo_ring.h, ), dtype=bool)
filling_pattern[0:n_bunch] = True
- beam = generate_beam(filling_pattern=filling_pattern, current_per_bunch=0.002)
+ beam = generate_beam(filling_pattern=filling_pattern,
+ current_per_bunch=0.002)
assert beam.current == pytest.approx(0.002 * n_bunch)
- assert beam.charge == pytest.approx(np.sum([bunch.charge for bunch in beam]))
- assert beam.particle_number == pytest.approx(np.sum([bunch.particle_number for bunch in beam]))
+ assert beam.charge == pytest.approx(
+ np.sum([bunch.charge for bunch in beam]))
+ assert beam.particle_number == pytest.approx(
+ np.sum([bunch.particle_number for bunch in beam]))
def test_save_and_load_beam_data(self, tmp_path, beam_uniform, demo_ring):
file_path = tmp_path / "beam_data"
beam_uniform.save(str(file_path))
loaded_beam = Beam(demo_ring)
loaded_beam.load(str(file_path) + ".hdf5", mpi=False)
- assert np.array_equal(beam_uniform.bunch_current, loaded_beam.bunch_current)
+ assert np.array_equal(beam_uniform.bunch_current,
+ loaded_beam.bunch_current)
assert np.array_equal(beam_uniform.bunch_mean, loaded_beam.bunch_mean)
assert np.array_equal(beam_uniform.bunch_std, loaded_beam.bunch_std)
- @pytest.mark.parametrize("var,option",
- [("bunch_current", None),
- ("bunch_charge", None),
- ("bunch_particle", None),
- ("bunch_mean","x"),
- ("bunch_std","x"),
- ("bunch_emit","x")])
- def test_plot_variables_with_respect_to_bunch_number(self, beam_uniform, var, option):
+ @pytest.mark.parametrize("var,option", [("bunch_current", None),
+ ("bunch_charge", None),
+ ("bunch_particle", None),
+ ("bunch_mean", "x"),
+ ("bunch_std", "x"),
+ ("bunch_emit", "x")])
+ def test_plot_variables_with_respect_to_bunch_number(
+ self, beam_uniform, var, option):
fig = beam_uniform.plot(var, option)
assert fig is not None
plt.close("all")
- def test_initialize_beam_mismatched_bunch_list_length(self, demo_ring, generate_bunch):
- mismatched_bunch_list = [generate_bunch() for _ in range(demo_ring.h - 1)]
+ def test_initialize_beam_mismatched_bunch_list_length(
+ self, demo_ring, generate_bunch):
+ mismatched_bunch_list = [
+ generate_bunch() for _ in range(demo_ring.h - 1)
+ ]
with pytest.raises(ValueError):
Beam(demo_ring, mismatched_bunch_list)
- def test_filling_pattern_and_distance_between_bunches(self, generate_beam, demo_ring):
- filling_pattern = np.ones((demo_ring.h,), dtype=bool)
+ def test_filling_pattern_and_distance_between_bunches(
+ self, generate_beam, demo_ring):
+ filling_pattern = np.ones((demo_ring.h, ), dtype=bool)
filling_pattern[5] = False
filling_pattern[8:10] = False
beam = generate_beam(filling_pattern=filling_pattern)
np.testing.assert_array_equal(beam.filling_pattern, filling_pattern)
- expected_distances = np.ones((demo_ring.h,))
+ expected_distances = np.ones((demo_ring.h, ))
expected_distances[4] = 2
expected_distances[5] = 0
expected_distances[7] = 3
expected_distances[8:10] = 0
- np.testing.assert_array_equal(beam.distance_between_bunches, expected_distances)
+ np.testing.assert_array_equal(beam.distance_between_bunches,
+ expected_distances)
- def test_update_filling_pattern_and_distance_between_bunches(self, beam_uniform, demo_ring):
- for i in [5,8,9]:
+ def test_update_filling_pattern_and_distance_between_bunches(
+ self, beam_uniform, demo_ring):
+ for i in [5, 8, 9]:
beam_uniform[i].charge = 0
beam_uniform[5].charge = 0
beam_uniform.update_filling_pattern()
beam_uniform.update_distance_between_bunches()
- expected_filling_pattern = np.ones((demo_ring.h,), dtype=bool)
+ expected_filling_pattern = np.ones((demo_ring.h, ), dtype=bool)
expected_filling_pattern[5] = False
expected_filling_pattern[8:10] = False
- np.testing.assert_array_equal(beam_uniform.filling_pattern, expected_filling_pattern)
+ np.testing.assert_array_equal(beam_uniform.filling_pattern,
+ expected_filling_pattern)
- expected_distances = np.ones((demo_ring.h,))
+ expected_distances = np.ones((demo_ring.h, ))
expected_distances[4] = 2
expected_distances[5] = 0
expected_distances[7] = 3
expected_distances[8:10] = 0
- np.testing.assert_array_equal(beam_uniform.distance_between_bunches, expected_distances)
+ np.testing.assert_array_equal(beam_uniform.distance_between_bunches,
+ expected_distances)
- def test_mpi_gather_and_close_consistency(self, mocker, demo_ring, generate_bunch):
+ def test_mpi_gather_and_close_consistency(self, mocker, demo_ring,
+ generate_bunch):
mock_mpi = mocker.patch('mbtrack2.tracking.parallel.Mpi')
mock_mpi_instance = mock_mpi.return_value
- mock_mpi_instance.comm.allgather.return_value = [generate_bunch() for _ in range(demo_ring.h)]
+ mock_mpi_instance.comm.allgather.return_value = [
+ generate_bunch() for _ in range(demo_ring.h)
+ ]
beam = Beam(ring=demo_ring)
beam.mpi_init()
beam.mpi_gather()
@@ -305,24 +397,26 @@ class TestBeam:
beam.mpi_close()
assert not beam.mpi_switch
assert beam.mpi is None
-
+
def test_init_bunch_list(self, demo_ring):
- filling_pattern = np.ones((demo_ring.h,), dtype=bool)
+ filling_pattern = np.ones((demo_ring.h, ), dtype=bool)
filling_pattern[5] = False
filling_pattern[8:10] = False
- bunch_list = [Bunch(demo_ring, mp_number=1, alive=filling_pattern[i]) for i in range(demo_ring.h)]
+ bunch_list = [
+ Bunch(demo_ring, mp_number=1, alive=filling_pattern[i])
+ for i in range(demo_ring.h)
+ ]
beam = Beam(demo_ring, bunch_list)
assert len(beam) == filling_pattern.sum()
np.testing.assert_array_equal(beam.filling_pattern, filling_pattern)
assert beam.distance_between_bunches is not None
-
+
def test_init_bunch_list_mpi(self, demo_ring, generate_bunch):
- filling_pattern = np.zeros((demo_ring.h,), dtype=bool)
+ filling_pattern = np.zeros((demo_ring.h, ), dtype=bool)
filling_pattern[0] = True
beam = Beam(demo_ring)
beam.init_bunch_list_mpi(generate_bunch(), filling_pattern)
-
+
assert len(beam) == 1
np.testing.assert_array_equal(beam.filling_pattern, filling_pattern)
assert beam.distance_between_bunches[0] == demo_ring.h
-
\ No newline at end of file
diff --git a/tests/unit/tracking/test_rf.py b/tests/unit/tracking/test_rf.py
index 6c0cb0e75190c64b9a7bd1776b620b32f2f0a6cf..6e052619e335d831aa5ceff3f9ecc783b183e1d1 100644
--- a/tests/unit/tracking/test_rf.py
+++ b/tests/unit/tracking/test_rf.py
@@ -1,33 +1,47 @@
import pytest
+
pytestmark = pytest.mark.unit
import numpy as np
from utility_test_functions import assert_attr_changed
-from mbtrack2 import (RFCavity,
- CavityResonator,
- ProportionalLoop,
- ProportionalIntegralLoop,
- TunerLoop,
- DirectFeedback)
+
+from mbtrack2 import (
+ CavityResonator,
+ DirectFeedback,
+ ProportionalIntegralLoop,
+ ProportionalLoop,
+ RFCavity,
+ TunerLoop,
+)
+
class TestRFCavity:
# Initialize RFCavity with valid parameters and verify attributes are changed when tracked
def test_rf_cavity(self, demo_ring, small_bunch):
- cavity = RFCavity(ring=demo_ring, m=1, Vc=1e6, theta=np.pi/4)
+ cavity = RFCavity(ring=demo_ring, m=1, Vc=1e6, theta=np.pi / 4)
assert_attr_changed(cavity, small_bunch, attrs_changed=["delta"])
+
@pytest.fixture
def cav_res(demo_ring):
- cav_res = CavityResonator(demo_ring, m=1, Rs=1e6, Q=1e4, QL=1e3, detune=-20e3, Ncav=4)
+ cav_res = CavityResonator(demo_ring,
+ m=1,
+ Rs=1e6,
+ Q=1e4,
+ QL=1e3,
+ detune=-20e3,
+ Ncav=4)
return cav_res
+
@pytest.fixture
def cav_res_tracking(demo_ring, cav_res):
cav_res.Vc = 1e6
- cav_res.theta = np.arccos(demo_ring.U0/cav_res.Vc)
- cav_res.set_generator(0.5)
+ cav_res.theta = np.arccos(demo_ring.U0 / cav_res.Vc)
+ cav_res.set_generator(0.5)
return cav_res
+
class TestCavityResonator:
@pytest.mark.parametrize("n_bin", [(75), (1)])
@@ -37,87 +51,102 @@ class TestCavityResonator:
initial_beam_phasor_record = cav_res_tracking.beam_phasor_record.copy()
beam_1bunch_mpi.mpi.share_distributions(beam_1bunch_mpi, n_bin=n_bin)
- assert_attr_changed(cav_res_tracking, beam_1bunch_mpi,attrs_changed=["delta"])
-
+ assert_attr_changed(cav_res_tracking,
+ beam_1bunch_mpi,
+ attrs_changed=["delta"])
+
assert not np.array_equal(cav_res_tracking.beam_phasor, initial_phasor)
- assert not np.array_equal(cav_res_tracking.beam_phasor_record, initial_beam_phasor_record)
+ assert not np.array_equal(cav_res_tracking.beam_phasor_record,
+ initial_beam_phasor_record)
@pytest.mark.parametrize("n_bin", [(75), (1)])
def test_track_no_mpi(self, cav_res_tracking, beam_uniform, n_bin):
cav_res_tracking.n_bin = n_bin
initial_phasor = cav_res_tracking.beam_phasor
initial_beam_phasor_record = cav_res_tracking.beam_phasor_record.copy()
-
- assert_attr_changed(cav_res_tracking, beam_uniform, attrs_changed=["delta"])
+
+ assert_attr_changed(cav_res_tracking,
+ beam_uniform,
+ attrs_changed=["delta"])
assert not np.array_equal(cav_res_tracking.beam_phasor, initial_phasor)
- assert not np.array_equal(cav_res_tracking.beam_phasor_record, initial_beam_phasor_record)
-
+ assert not np.array_equal(cav_res_tracking.beam_phasor_record,
+ initial_beam_phasor_record)
+
@pytest.mark.parametrize("n_bin", [(75), (1)])
- def test_track_no_mpi_non_uniform(self, cav_res_tracking, beam_non_uniform, n_bin):
+ def test_track_no_mpi_non_uniform(self, cav_res_tracking, beam_non_uniform,
+ n_bin):
cav_res_tracking.n_bin = n_bin
initial_phasor = cav_res_tracking.beam_phasor
initial_beam_phasor_record = cav_res_tracking.beam_phasor_record.copy()
-
- assert_attr_changed(cav_res_tracking, beam_non_uniform, attrs_changed=["delta"])
+
+ assert_attr_changed(cav_res_tracking,
+ beam_non_uniform,
+ attrs_changed=["delta"])
assert not np.array_equal(cav_res_tracking.beam_phasor, initial_phasor)
- assert not np.array_equal(cav_res_tracking.beam_phasor_record, initial_beam_phasor_record)
-
+ assert not np.array_equal(cav_res_tracking.beam_phasor_record,
+ initial_beam_phasor_record)
+
# Track a beam with zero charge per macro-particle
def test_track_zero_charge_per_mp(self, cav_res_tracking, generate_beam):
beam = generate_beam(current_per_bunch=0)
assert_attr_changed(cav_res_tracking, beam, change=False)
-
+
# Apply RF feedback loops correctly during tracking
def test_apply_rf_feedback(self, cav_res_tracking, beam_uniform, mocker):
cav_res_tracking.feedback.append(mocker.Mock())
cav_res_tracking.track(beam_uniform)
- cav_res_tracking.feedback[0].track.assert_called_once(), "RF feedback should be applied"
-
+ cav_res_tracking.feedback[0].track.assert_called_once(
+ ), "RF feedback should be applied"
+
@pytest.mark.parametrize("ref_frame", [("beam"), ("rf")])
def test_phasor_decay(self, cav_res, ref_frame):
init_phasor = 1 + 1j
cav_res.beam_phasor = init_phasor
cav_res.phasor_decay(10e-6, ref_frame)
assert cav_res.beam_phasor != init_phasor
-
+
@pytest.mark.parametrize("ref_frame", [("beam"), ("rf")])
def test_phasor_evol(self, cav_res, ref_frame):
init_phasor = 1 + 1j
cav_res.beam_phasor = init_phasor
- profile = np.random.randint(1000, size=(10,))
+ profile = np.random.randint(1000, size=(10, ))
bin_length = 1e-12
charge_per_mp = 1e-15
cav_res.phasor_evol(profile, bin_length, charge_per_mp, ref_frame)
assert cav_res.beam_phasor != init_phasor
-
+
def test_phasor_init(self, cav_res, beam_non_uniform):
init_phasor = cav_res.beam_phasor
cav_res.init_phasor_track(beam_non_uniform)
phasor_init_phasor_track = cav_res.beam_phasor
-
+
cav_res.beam_phasor = init_phasor
cav_res.init_phasor(beam_non_uniform)
phasor_init_phasor = cav_res.beam_phasor
-
+
assert phasor_init_phasor != init_phasor
assert phasor_init_phasor_track != init_phasor
- assert np.allclose(phasor_init_phasor, phasor_init_phasor_track, rtol=1e-2)
-
+ assert np.allclose(phasor_init_phasor,
+ phasor_init_phasor_track,
+ rtol=1e-2)
+
def test_phasor_init_mpi(self, cav_res, beam_1bunch_mpi):
init_phasor = cav_res.beam_phasor
cav_res.init_phasor_track(beam_1bunch_mpi)
phasor_init_phasor_track = cav_res.beam_phasor
-
+
cav_res.beam_phasor = init_phasor
cav_res.init_phasor(beam_1bunch_mpi)
phasor_init_phasor = cav_res.beam_phasor
-
+
assert phasor_init_phasor != init_phasor
assert phasor_init_phasor_track != init_phasor
- assert np.allclose(phasor_init_phasor, phasor_init_phasor_track, rtol=1e-2)
-
+ assert np.allclose(phasor_init_phasor,
+ phasor_init_phasor_track,
+ rtol=1e-2)
+
# Setting detune updates _detune, _fr, _wr, and _psi correctly
def test_detune(self, cav_res):
detune_value = 1000
@@ -125,18 +154,20 @@ class TestCavityResonator:
assert cav_res._detune == detune_value
assert cav_res._fr == detune_value + cav_res.m * cav_res.ring.f1
assert cav_res._wr == cav_res._fr * 2 * np.pi
- assert cav_res._psi == np.arctan(cav_res.QL * (cav_res._fr / (cav_res.m * cav_res.ring.f1) - (cav_res.m * cav_res.ring.f1) / cav_res._fr))
+ assert cav_res._psi == np.arctan(
+ cav_res.QL * (cav_res._fr / (cav_res.m * cav_res.ring.f1) -
+ (cav_res.m * cav_res.ring.f1) / cav_res._fr))
# update_feedback is called after setting detune
def test_update_feedback_called(self, mocker, cav_res):
mocker.spy(cav_res, 'update_feedback')
cav_res.detune = 1000
- cav_res.update_feedback.assert_called_once()
+ cav_res.update_feedback.assert_called_once()
def test_plot_phasor_diagram(self, cav_res_tracking):
fig = cav_res_tracking.plot_phasor(0.5)
assert fig is not None
-
+
# Calculating generator phasor using Vg and theta_g
def test_generator_phasor_calculation(self, cav_res):
cav_res.Vg = 1000
@@ -154,15 +185,16 @@ class TestCavityResonator:
# Retrieving cavity phasor record for each bunch
def test_cavity_phasor_record_retrieval(self, cav_res):
- cav_res.generator_phasor_record = np.array([1+1j, 2+2j])
- cav_res.beam_phasor_record = np.array([0.5+0.5j, 1+1j])
- expected_record = np.array([1.5+1.5j, 3+3j])
+ cav_res.generator_phasor_record = np.array([1 + 1j, 2 + 2j])
+ cav_res.beam_phasor_record = np.array([0.5 + 0.5j, 1 + 1j])
+ expected_record = np.array([1.5 + 1.5j, 3 + 3j])
assert np.allclose(cav_res.cavity_phasor_record, expected_record)
# Accessing ig phasor record when feedback is present
def test_ig_phasor_record_with_feedback(self, cav_res):
- pi_loop = ProportionalIntegralLoop(cav_res.ring, cav_res, [1.0, 1.0], 10, 10, 10)
- pi_loop.ig_phasor_record = np.array([1+1j, 2+2j])
+ pi_loop = ProportionalIntegralLoop(cav_res.ring, cav_res, [1.0, 1.0],
+ 10, 10, 10)
+ pi_loop.ig_phasor_record = np.array([1 + 1j, 2 + 2j])
cav_res.feedback.append(pi_loop)
assert np.allclose(cav_res.ig_phasor_record, pi_loop.ig_phasor_record)
@@ -195,7 +227,7 @@ class TestCavityResonator:
# Set up the initial conditions
cav_res.Rs = 1e6 # Shunt impedance
cav_res.QL = 1e3 # Loaded quality factor
- cav_res.Q = 1e4 # Quality factor
+ cav_res.Q = 1e4 # Quality factor
# Calculate expected RL
expected_RL = cav_res.Rs / (1 + (cav_res.Q / cav_res.QL - 1))
@@ -211,26 +243,27 @@ class TestCavityResonator:
cav_res.set_optimal_detune(0.5)
assert cav_res.beta != initial_beta
assert cav_res.psi != initial_psi
-
+
# Verify that is_CBI_stable correctly calculates growth rate and mode number
def test_calculate_growth_rate_and_mode(self, cav_res_tracking):
growth_rate, mu = cav_res_tracking.is_CBI_stable(I0=0.5)
assert isinstance(growth_rate, float)
assert np.issubdtype(type(mu), np.integer)
-
+
# Ensure is_CBI_stable returns growth rates for specified modes when modes is provided
def test_return_growth_rates_for_specified_modes(self, cav_res_tracking):
modes = [0, 1, 2]
growth_rates = cav_res_tracking.is_CBI_stable(I0=0.5, modes=modes)
assert len(growth_rates) == len(modes)
-
+
def test_is_DC_Robinson_stable(self, cav_res_tracking):
- assert isinstance(cav_res_tracking.is_DC_Robinson_stable(0.5), (bool,np.bool_))
-
+ assert isinstance(cav_res_tracking.is_DC_Robinson_stable(0.5),
+ (bool, np.bool_))
+
def test_plot_DC_Robinson_stability(self, cav_res_tracking):
fig = cav_res_tracking.plot_DC_Robinson_stability()
assert fig is not None
-
+
@pytest.mark.parametrize("z,expected_type",
[(1e-3, float),
(np.linspace(-1e-3, 1e-3, 100), np.ndarray)])
@@ -239,20 +272,21 @@ class TestCavityResonator:
assert isinstance(cav_res_tracking.dVRF(z, 0.5), expected_type)
assert isinstance(cav_res_tracking.ddVRF(z, 0.5), expected_type)
assert isinstance(cav_res_tracking.deltaVRF(z, 0.5), expected_type)
-
- def test_sample_voltage_default_parameters(self, cav_res_tracking, beam_uniform):
+
+ def test_sample_voltage_default_parameters(self, cav_res_tracking,
+ beam_uniform):
cav_res_tracking.init_tracking(beam_uniform)
- n_points=1e4
+ n_points = 1e4
beam_phasor_init = 1 + 1j
cav_res_tracking.beam_phasor = beam_phasor_init
pos, voltage_rec = cav_res_tracking.sample_voltage(n_points=n_points)
assert len(pos) == n_points
assert len(voltage_rec) == n_points
assert np.allclose(cav_res_tracking.beam_phasor, beam_phasor_init)
-
+
def test_set_generator(self, cav_res, demo_ring):
cav_res.Vc = 1e6
- cav_res.theta = np.arccos(demo_ring.U0/cav_res.Vc)
+ cav_res.theta = np.arccos(demo_ring.U0 / cav_res.Vc)
cav_res.set_generator(0.5)
assert cav_res.Pg is not None
assert cav_res.Vgr is not None
@@ -260,24 +294,29 @@ class TestCavityResonator:
assert cav_res.Vg is not None
assert cav_res.theta_g is not None
assert cav_res.generator_phasor_record is not None
-
+
def test_functions(self, cav_res_tracking):
assert isinstance(cav_res_tracking.Pb(0.5), float)
assert isinstance(cav_res_tracking.Pr(0.5), float)
assert isinstance(cav_res_tracking.Vbr(0.5), float)
assert isinstance(cav_res_tracking.Vb(0.5), float)
assert isinstance(cav_res_tracking.Z(1e9), complex)
-
+
def test_properties(self, cav_res_tracking):
assert cav_res_tracking.Pc > 0
assert cav_res_tracking.filling_time > 0
assert cav_res_tracking.loss_factor > 0
-
+
+
class TestProportionalLoop:
-
+
@pytest.fixture
def prop_loop(self, demo_ring, cav_res_tracking):
- return ProportionalLoop(demo_ring, cav_res_tracking, gain_A=0.5, gain_P=0.5, delay=2)
+ return ProportionalLoop(demo_ring,
+ cav_res_tracking,
+ gain_A=0.5,
+ gain_P=0.5,
+ delay=2)
# track method updates cav_res.Vg and cav_res.theta_g correctly
def test_track_delay(self, prop_loop):
@@ -286,13 +325,13 @@ class TestProportionalLoop:
prop_loop.track()
assert prop_loop.cav_res.Vg == initial_Vg
assert prop_loop.cav_res.theta_g == initial_theta_g
-
+
# track method updates cav_res.Vg and cav_res.theta_g correctly
def test_track_update(self, prop_loop):
initial_Vg = prop_loop.cav_res.Vg
initial_theta_g = prop_loop.cav_res.theta_g
prop_loop.track()
- prop_loop.track() # delay=2 is 3 turns to take action
+ prop_loop.track() # delay=2 is 3 turns to take action
prop_loop.track()
assert prop_loop.cav_res.Vg != initial_Vg
assert prop_loop.cav_res.theta_g != initial_theta_g
@@ -300,19 +339,32 @@ class TestProportionalLoop:
# Initialization with delay less than 1 raises ValueError
def test_initialization_delay_less_than_one(self, demo_ring, cav_res):
with pytest.raises(ValueError):
- ProportionalLoop(demo_ring, cav_res, gain_A=0.5, gain_P=0.5, delay=0)
+ ProportionalLoop(demo_ring,
+ cav_res,
+ gain_A=0.5,
+ gain_P=0.5,
+ delay=0)
# Initialization with non-integer delay converts it to integer
def test_initialization_non_integer_delay(self, demo_ring, cav_res):
- loop = ProportionalLoop(demo_ring, cav_res, gain_A=0.5, gain_P=0.5, delay=2.7)
+ loop = ProportionalLoop(demo_ring,
+ cav_res,
+ gain_A=0.5,
+ gain_P=0.5,
+ delay=2.7)
assert loop.delay == 2
-
+
+
class TestTunerLoop:
-
+
@pytest.fixture
def tuner_loop(self, demo_ring, cav_res_tracking):
- return TunerLoop(demo_ring, cav_res_tracking, gain=1, avering_period=2, offset=0)
-
+ return TunerLoop(demo_ring,
+ cav_res_tracking,
+ gain=1,
+ avering_period=2,
+ offset=0)
+
def test_track_once(self, tuner_loop):
initial_psi = tuner_loop.cav_res.psi
tuner_loop.track()
@@ -323,14 +375,17 @@ class TestTunerLoop:
# Track method correctly adjusts cav_res.psi when count equals avering_period
def test_track_adjusts_psi(self, tuner_loop):
initial_psi = tuner_loop.cav_res.psi
- for _ in range(tuner_loop.avering_period+1):
+ for _ in range(tuner_loop.avering_period + 1):
tuner_loop.track()
assert tuner_loop.diff == 0
assert tuner_loop.cav_res.psi != initial_psi
# Initialize with avering_period set to None and verify default calculation
- def test_default_averaging_period_calculation(self, demo_ring, cav_res_tracking):
- tuner_loop = TunerLoop(demo_ring, cav_res_tracking, avering_period=None)
+ def test_default_averaging_period_calculation(self, demo_ring,
+ cav_res_tracking):
+ tuner_loop = TunerLoop(demo_ring,
+ cav_res_tracking,
+ avering_period=None)
fs = demo_ring.synchrotron_tune(cav_res_tracking.Vc) * demo_ring.f0
expected_period = int(2 / fs / demo_ring.T0)
assert tuner_loop.avering_period == expected_period
@@ -343,27 +398,33 @@ class TestTunerLoop:
tuner_loop.track()
assert tuner_loop.cav_res.psi == initial_psi
+
class TestProportionalIntegralLoop:
-
+
@pytest.fixture
def pi_loop(self, demo_ring, cav_res_tracking):
- loop = ProportionalIntegralLoop(demo_ring, cav_res_tracking, [0.5, 1e4], 8, 7, 5)
+ loop = ProportionalIntegralLoop(demo_ring, cav_res_tracking,
+ [0.5, 1e4], 8, 7, 5)
return loop
def test_track(self, pi_loop):
initial_ig_phasor = pi_loop.ig_phasor.copy()
- generator_phasor_record_init = pi_loop.cav_res.generator_phasor_record.copy()
+ generator_phasor_record_init = pi_loop.cav_res.generator_phasor_record.copy(
+ )
pi_loop.track()
assert not np.array_equal(initial_ig_phasor, pi_loop.ig_phasor)
- assert not np.array_equal(generator_phasor_record_init, pi_loop.cav_res.generator_phasor_record)
-
+ assert not np.array_equal(generator_phasor_record_init,
+ pi_loop.cav_res.generator_phasor_record)
+
# Ig2Vg method correctly updates generator phasor record and cavity parameters
def test_Ig2Vg_updates_generator_phasor(self, pi_loop):
- generator_phasor_record_init = pi_loop.cav_res.generator_phasor_record.copy()
+ generator_phasor_record_init = pi_loop.cav_res.generator_phasor_record.copy(
+ )
Vg_init = pi_loop.cav_res.Vg
theta_g_init = pi_loop.cav_res.theta_g
pi_loop.Ig2Vg()
- assert not np.allclose(generator_phasor_record_init, pi_loop.cav_res.generator_phasor_record)
+ assert not np.allclose(generator_phasor_record_init,
+ pi_loop.cav_res.generator_phasor_record)
assert pi_loop.cav_res.Vg != Vg_init
assert pi_loop.cav_res.theta_g != theta_g_init
@@ -377,27 +438,28 @@ class TestProportionalIntegralLoop:
def test_init_FFconst_with_FF_true(self, pi_loop):
pi_loop.init_FFconst()
assert pi_loop.FFconst != 0
-
+
def test_parameters(self, pi_loop):
assert isinstance(pi_loop.Vg2Ig(1e6), complex)
pi_loop.IIR_init(1e9)
assert isinstance(pi_loop.IIRcutoff, float)
assert isinstance(pi_loop.IIR(1.0 + 1.0j), complex)
+
class TestDirectFeedback:
-
+
@pytest.fixture
def drf_fb(self, demo_ring, cav_res_tracking):
- drf_fb = DirectFeedback(DFB_gain=1.0,
- DFB_phase_shift=0.1,
- ring=demo_ring,
- cav_res=cav_res_tracking,
- gain=[0.5, 1e4],
- sample_num=8,
- every=7,
- delay=5)
+ drf_fb = DirectFeedback(DFB_gain=1.0,
+ DFB_phase_shift=0.1,
+ ring=demo_ring,
+ cav_res=cav_res_tracking,
+ gain=[0.5, 1e4],
+ sample_num=8,
+ every=7,
+ delay=5)
return drf_fb
-
+
def test_properties(self, drf_fb):
assert isinstance(drf_fb.DFB_phase_shift, float)
assert isinstance(drf_fb.phase_shift, float)
@@ -405,7 +467,7 @@ class TestDirectFeedback:
assert isinstance(drf_fb.DFB_alpha, float)
assert isinstance(drf_fb.DFB_gamma, float)
assert isinstance(drf_fb.DFB_Rs, float)
-
+
def test_methods(self, drf_fb):
vg_main, vg_drf = drf_fb.DFB_Vg()
assert isinstance(vg_main, complex)
@@ -415,20 +477,23 @@ class TestDirectFeedback:
def test_track(self, drf_fb):
initial_ig_phasor_record = drf_fb.ig_phasor_record.copy()
initial_dfb_ig_phasor = drf_fb.DFB_ig_phasor.copy()
- generator_phasor_record_init = drf_fb.cav_res.generator_phasor_record.copy()
+ generator_phasor_record_init = drf_fb.cav_res.generator_phasor_record.copy(
+ )
drf_fb.track()
- assert not np.array_equal(drf_fb.ig_phasor_record, initial_ig_phasor_record)
+ assert not np.array_equal(drf_fb.ig_phasor_record,
+ initial_ig_phasor_record)
assert not np.array_equal(drf_fb.DFB_ig_phasor, initial_dfb_ig_phasor)
- assert not np.array_equal(generator_phasor_record_init, drf_fb.cav_res.generator_phasor_record)
+ assert not np.array_equal(generator_phasor_record_init,
+ drf_fb.cav_res.generator_phasor_record)
# Set DFB parameters using DFB_parameter_set and verify changes in DFB_ig_phasor
def test_set_dfb_parameters(self, drf_fb):
initial_DFB_ig_phasor = drf_fb.DFB_ig_phasor.copy()
initial_ig_phasor = drf_fb.ig_phasor.copy()
-
+
drf_fb.DFB_parameter_set(DFB_gain=2.0, DFB_phase_shift=0.2)
-
+
assert not np.array_equal(drf_fb.DFB_ig_phasor, initial_DFB_ig_phasor)
assert not np.array_equal(drf_fb.ig_phasor, initial_ig_phasor)
assert drf_fb.DFB_gain == 2.0
- assert drf_fb.DFB_phase_shift == 0.2
\ No newline at end of file
+ assert drf_fb.DFB_phase_shift == 0.2
diff --git a/tests/unit/tracking/test_spacecharge.py b/tests/unit/tracking/test_spacecharge.py
index 20d3073fad389887d1cab0e962287fe4511f01c9..6f42e4cf7d2167cb7b30aad6b83fbe7c94aa4254 100644
--- a/tests/unit/tracking/test_spacecharge.py
+++ b/tests/unit/tracking/test_spacecharge.py
@@ -1,8 +1,11 @@
import pytest
+
pytestmark = pytest.mark.unit
-from mbtrack2 import TransverseSpaceCharge
-from utility_test_functions import assert_attr_changed
import numpy as np
+from utility_test_functions import assert_attr_changed
+
+from mbtrack2 import TransverseSpaceCharge
+
class TestTransverseSpaceCharge:
@@ -11,7 +14,7 @@ class TestTransverseSpaceCharge:
tsc = TransverseSpaceCharge(demo_ring, 1000.0)
large_bunch.alive[5:700] = False
assert_attr_changed(tsc, large_bunch, attrs_changed=["xp", "yp"])
-
+
# Track a bunch with valid x, y coordinates through the space charge element
def test_track_without_track_alive(self, demo_ring, large_bunch):
tsc = TransverseSpaceCharge(demo_ring, 1000.0)
@@ -21,7 +24,10 @@ class TestTransverseSpaceCharge:
# Initialize with zero interaction length and verify no kicks are applied
def test_zero_interaction_length_no_kicks(self, demo_ring, small_bunch):
tsc = TransverseSpaceCharge(demo_ring, 0.0)
- assert_attr_changed(tsc, small_bunch, attrs_changed=["xp", "yp"], change=False)
+ assert_attr_changed(tsc,
+ small_bunch,
+ attrs_changed=["xp", "yp"],
+ change=False)
# Handle a bunch with bins having zero particles without errors
def test_handle_bins_with_zero_particles(self, demo_ring, small_bunch):
@@ -31,8 +37,8 @@ class TestTransverseSpaceCharge:
assert True
# Test with different n_bins
- @pytest.mark.parametrize("n_bins",[(1),(2),(100)])
+ @pytest.mark.parametrize("n_bins", [(1), (2), (100)])
def test_n_bins_zero_behavior(self, demo_ring, small_bunch, n_bins):
tsc = TransverseSpaceCharge(demo_ring, 1.0, n_bins=n_bins)
tsc.track(small_bunch)
- assert True
\ No newline at end of file
+ assert True
diff --git a/tests/unit/tracking/test_synchrotron.py b/tests/unit/tracking/test_synchrotron.py
index daf706ca590a2c2a81bed09572d102d2460b58d3..b75280bc53f37b70d21bda15f27aecdc64931891 100644
--- a/tests/unit/tracking/test_synchrotron.py
+++ b/tests/unit/tracking/test_synchrotron.py
@@ -1,10 +1,13 @@
import numpy as np
import pytest
+
pytestmark = pytest.mark.unit
import at
from scipy.constants import c, e
+
from mbtrack2 import Electron, Synchrotron
+
@pytest.fixture
def demo_ring(local_optics):
h = 20
@@ -15,42 +18,64 @@ def demo_ring(local_optics):
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])
+ 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)
-
+ 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 demo_ring_h1(demo_ring):
demo_ring.h = 1
return demo_ring
+
@pytest.fixture
def generate_ring_with_at_lattice(generate_at_optics):
+
def generate(**kwargs):
optics = generate_at_optics(**kwargs)
h = 416
tau = np.array([6.56e-3, 6.56e-3, 3.27e-3])
- emit = np.array([3.9e-9, 3.9e-9*0.01])
+ emit = np.array([3.9e-9, 3.9e-9 * 0.01])
sigma_0 = 15e-12
sigma_delta = 1.025e-3
particle = Electron()
- ring = Synchrotron(h, optics, particle, tau=tau, emit=emit,
- sigma_0=sigma_0, sigma_delta=sigma_delta, **kwargs)
+ ring = Synchrotron(h,
+ optics,
+ particle,
+ tau=tau,
+ emit=emit,
+ sigma_0=sigma_0,
+ sigma_delta=sigma_delta,
+ **kwargs)
return ring
+
return generate
+
@pytest.fixture
def ring_with_at_lattice(generate_ring_with_at_lattice):
ring = generate_ring_with_at_lattice()
return ring
+
class TestSynchrotron:
def test_synchrotron_values(self, demo_ring):
@@ -62,11 +87,11 @@ class TestSynchrotron:
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])
+ emit = np.array([50e-9, 50e-9 * 0.01])
sigma_0 = 30e-12
sigma_delta = 1e-3
- chro = [1.0,1.0]
-
+ chro = [1.0, 1.0]
+
assert pytest.approx(demo_ring.h) == h
assert pytest.approx(demo_ring.L) == L
assert pytest.approx(demo_ring.E0) == E0
@@ -78,56 +103,67 @@ class TestSynchrotron:
assert pytest.approx(demo_ring.sigma_0) == sigma_0
assert pytest.approx(demo_ring.sigma_delta) == sigma_delta
np.testing.assert_allclose(demo_ring.chro, chro)
- assert pytest.approx(demo_ring.T0) == L/c
- assert pytest.approx(demo_ring.T1) == L/c/h
- assert pytest.approx(demo_ring.f0) == c/L
- assert pytest.approx(demo_ring.f1) == 1/(L/c/h)
- assert pytest.approx(demo_ring.omega0) == 2 * np.pi * c/L
- assert pytest.approx(demo_ring.omega1) == 2 * np.pi * 1/(L/c/h)
- assert pytest.approx(demo_ring.k1) == 2 * np.pi * 1/(L/c/h) / c
- assert pytest.approx(demo_ring.gamma) == E0 / (particle.mass * c**2 / e)
- assert pytest.approx(demo_ring.beta) == np.sqrt(1 - (E0 / (particle.mass * c**2 / e))**-2)
-
+ assert pytest.approx(demo_ring.T0) == L / c
+ assert pytest.approx(demo_ring.T1) == L / c / h
+ assert pytest.approx(demo_ring.f0) == c / L
+ assert pytest.approx(demo_ring.f1) == 1 / (L/c/h)
+ assert pytest.approx(demo_ring.omega0) == 2 * np.pi * c / L
+ assert pytest.approx(demo_ring.omega1) == 2 * np.pi * 1 / (L/c/h)
+ assert pytest.approx(demo_ring.k1) == 2 * np.pi * 1 / (L/c/h) / c
+ assert pytest.approx(
+ demo_ring.gamma) == E0 / (particle.mass * c**2 / e)
+ assert pytest.approx(
+ demo_ring.beta) == np.sqrt(1 - (E0 /
+ (particle.mass * c**2 / e))**-2)
+
def test_synchrotron_mcf(self, 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)
-
+ 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(self, demo_ring):
tuneS = demo_ring.synchrotron_tune(1e6)
assert pytest.approx(tuneS, rel=1e-4) == 0.0017553
-
+
def test_synchrotron_sigma(self, demo_ring):
- np.testing.assert_allclose(demo_ring.sigma(), np.array([2.23606798e-04, 2.23606798e-04, 2.23606798e-05, 2.23606798e-05]))
+ np.testing.assert_allclose(
+ demo_ring.sigma(),
+ np.array([
+ 2.23606798e-04, 2.23606798e-04, 2.23606798e-05, 2.23606798e-05
+ ]))
def test_synchrotron_sigma_position(self, ring_with_at_lattice):
pos = np.linspace(0, ring_with_at_lattice.L, 100)
sig = ring_with_at_lattice.sigma(pos)
assert sig.shape == (4, 100)
-
+
def test_get_adts(self, ring_with_at_lattice):
ring_with_at_lattice.get_adts()
assert ring_with_at_lattice.adts is not None
-
+
def test_get_adts_plot(self, ring_with_at_lattice):
axes = ring_with_at_lattice.get_adts(plot=True)
assert axes is not None
-
+
def test_get_chroma(self, ring_with_at_lattice):
ring_with_at_lattice.get_chroma()
assert len(ring_with_at_lattice.chro) == 8
-
+
def test_get_mcf_order(self, ring_with_at_lattice):
ring_with_at_lattice.get_mcf_order()
assert len(ring_with_at_lattice.mcf_order) == 3
-
+
def test_synchrotron_long_twiss(self, 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)
+ 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_to_pyat(self, demo_ring):
pyat_simple_ring = demo_ring.to_pyat(1e6)
- assert isinstance(pyat_simple_ring, at.lattice.lattice_object.Lattice)
\ No newline at end of file
+ assert isinstance(pyat_simple_ring, at.lattice.lattice_object.Lattice)
diff --git a/tests/unit/tracking/test_wakepotential.py b/tests/unit/tracking/test_wakepotential.py
index 403f711d3331af1234000bae1b6ed85d1787c102..ed4327e615aa4e9167e71d9e0806cd1c904dd915 100644
--- a/tests/unit/tracking/test_wakepotential.py
+++ b/tests/unit/tracking/test_wakepotential.py
@@ -1,35 +1,44 @@
import pytest
+
pytestmark = pytest.mark.unit
import numpy as np
import pandas as pd
-from mbtrack2 import WakePotential, Resonator, LongRangeResistiveWall
from utility_test_functions import assert_attr_changed
+from mbtrack2 import LongRangeResistiveWall, Resonator, WakePotential
+
+
@pytest.fixture
def generate_resonator():
- def generate(time=np.linspace(-1000e-12, 1000e-12, int(1e5)),
- frequency=np.linspace(0.1e6, 100e9, int(1e5)),
- Rs=100,
- fr=1e9,
- Q=20,
+
+ def generate(time=np.linspace(-1000e-12, 1000e-12, int(1e5)),
+ frequency=np.linspace(0.1e6, 100e9, int(1e5)),
+ Rs=100,
+ fr=1e9,
+ Q=20,
plane="long"):
res = Resonator(time, frequency, Rs, fr, Q, plane)
return res
+
return generate
+
@pytest.fixture
def generate_wakepotential(demo_ring, generate_resonator):
- def generate(ring=demo_ring,
- wakefield=None,
- n_bin=80,
- interp_on_postion=True,
+
+ def generate(ring=demo_ring,
+ wakefield=None,
+ n_bin=80,
+ interp_on_postion=True,
**kwargs):
if wakefield is None:
wakefield = generate_resonator(**kwargs)
wp = WakePotential(ring, wakefield, n_bin, interp_on_postion)
return wp
+
return generate
+
class TestWakePotential:
# Compute charge density profile for a given Bunch object
@@ -41,15 +50,17 @@ class TestWakePotential:
# Calculate dipole moment for a Bunch object on specified plane
@pytest.mark.parametrize("plane", [("x"), ("y")])
- def test_dipole_moment_calculation(self, generate_wakepotential, small_bunch, plane):
+ def test_dipole_moment_calculation(self, generate_wakepotential,
+ small_bunch, plane):
wp = generate_wakepotential()
wp.charge_density(small_bunch)
dipole = wp.dipole_moment(small_bunch, plane, wp.tau)
assert len(dipole) == len(wp.tau)
# Prepare wake function for a specified wake type
- @pytest.mark.parametrize("plane", [("x"), ("y"),("long")])
- def test_prepare_wakefunction_smaller_window(self, generate_wakepotential, plane):
+ @pytest.mark.parametrize("plane", [("x"), ("y"), ("long")])
+ def test_prepare_wakefunction_smaller_window(self, generate_wakepotential,
+ plane):
wp = generate_wakepotential(plane=plane)
tau = np.linspace(-100e-12, 100e-12, int(1e5))
if plane == "x" or plane == "y":
@@ -59,10 +70,11 @@ class TestWakePotential:
tau0, dtau0, W0 = wp.prepare_wakefunction(comp, tau)
assert len(tau0) > 0
assert len(W0) > 0
-
+
# Prepare wake function for a specified wake type
- @pytest.mark.parametrize("plane", [("x"), ("y"),("long")])
- def test_prepare_wakefunction_larger_window(self, generate_wakepotential, plane):
+ @pytest.mark.parametrize("plane", [("x"), ("y"), ("long")])
+ def test_prepare_wakefunction_larger_window(self, generate_wakepotential,
+ plane):
wp = generate_wakepotential(plane=plane)
tau = np.linspace(-2000e-12, 2000e-12, int(1e5))
if plane == "x" or plane == "y":
@@ -74,8 +86,9 @@ class TestWakePotential:
assert len(W0) > 0
# Compute wake potential for a Bunch object and wake type
- @pytest.mark.parametrize("plane", [("x"), ("y"),("long")])
- def test_compute_wakepotential(self, generate_wakepotential, small_bunch, plane):
+ @pytest.mark.parametrize("plane", [("x"), ("y"), ("long")])
+ def test_compute_wakepotential(self, generate_wakepotential, small_bunch,
+ plane):
wp = generate_wakepotential(plane=plane)
wp.charge_density(small_bunch)
if plane == "x" or plane == "y":
@@ -86,23 +99,32 @@ class TestWakePotential:
assert len(Wp) == len(tau0)
# Track a Bunch object through the WakePotential element
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("long","delta")])
- def test_track_bunch(self, generate_wakepotential, small_bunch, plane, attr):
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("long", "delta")])
+ def test_track_bunch(self, generate_wakepotential, small_bunch, plane,
+ attr):
wp = generate_wakepotential(plane=plane)
assert_attr_changed(wp, small_bunch, attrs_changed=[attr])
-
+
# Track a Bunch object through the WakePotential element
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("long","delta")])
- def test_track_bunch_no_exact_interp(self, generate_wakepotential, small_bunch, plane, attr):
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("long", "delta")])
+ def test_track_bunch_no_exact_interp(self, generate_wakepotential,
+ small_bunch, plane, attr):
wp = generate_wakepotential(plane=plane, interp_on_postion=False)
assert_attr_changed(wp, small_bunch, attrs_changed=[attr])
# Handle empty Bunch object in track method
- @pytest.mark.parametrize("plane, attr", [("x","xp"), ("y","yp"),("long","delta")])
- def test_track_empty_bunch(self, generate_wakepotential, small_bunch, plane, attr):
+ @pytest.mark.parametrize("plane, attr", [("x", "xp"), ("y", "yp"),
+ ("long", "delta")])
+ def test_track_empty_bunch(self, generate_wakepotential, small_bunch,
+ plane, attr):
wp = generate_wakepotential(plane=plane)
small_bunch.alive[:] = False
- assert_attr_changed(wp, small_bunch, attrs_changed=[attr], change=False)
+ assert_attr_changed(wp,
+ small_bunch,
+ attrs_changed=[attr],
+ change=False)
# Manage non-uniformly sampled wake functions
def test_non_uniform_sampling_error(self, generate_wakepotential):
@@ -111,85 +133,90 @@ class TestWakePotential:
wp.check_sampling()
# Calculate reference loss factor and compare to Gaussian bunch
- @pytest.mark.parametrize("plane", [("x"), ("y"),("long")])
- def test_reference_loss_factor_comparison(self, generate_wakepotential, large_bunch, plane):
+ @pytest.mark.parametrize("plane", [("x"), ("y"), ("long")])
+ def test_reference_loss_factor_comparison(self, generate_wakepotential,
+ large_bunch, plane):
# Generate a WakePotential instance
wp = generate_wakepotential(plane=plane)
-
+
# Calculate the reference loss factor using the small bunch
large_bunch["x"] += 1e-3
large_bunch["y"] += 1e-3
wp.track(large_bunch)
loss_data = wp.reference_loss(large_bunch)
-
+
# Check if the loss data is a DataFrame and contains expected columns
- assert isinstance(loss_data, pd.DataFrame), "Loss data should be a DataFrame"
+ assert isinstance(loss_data,
+ pd.DataFrame), "Loss data should be a DataFrame"
assert 'TD factor' in loss_data.columns, "DataFrame should contain 'TD factor' column"
assert 'FD factor' in loss_data.columns, "DataFrame should contain 'FD factor' column"
assert 'Relative error [%]' in loss_data.columns, "DataFrame should contain 'Relative error [%]' column"
-
+
# Verify that the calculated loss factors are within a reasonable range
for index, row in loss_data.iterrows():
- assert abs(row['Relative error [%]']) < 1, f"Relative error for {index} is too high"
+ assert abs(row['Relative error [%]']
+ ) < 1, f"Relative error for {index} is too high"
# Reduce wake function sampling by a specified factor
def test_reduce_sampling(self, generate_wakepotential):
# Create a WakePotential instance with a mock wakefield
wp = generate_wakepotential()
-
+
# Original length of the wake function data
original_length = len(wp.wakefield.Wlong.data.index)
-
+
# Define a reduction factor
reduction_factor = 2
-
+
# Reduce the sampling of the wake function
wp.reduce_sampling(reduction_factor)
-
+
# New length of the wake function data after reduction
new_length = len(wp.wakefield.Wlong.data.index)
-
+
# Assert that the new length is approximately half of the original length
assert new_length == original_length // reduction_factor
# Test plotting functions with different plot options
- def test_plot_last_wake_with_various_options(self, generate_wakepotential, small_bunch):
+ def test_plot_last_wake_with_various_options(self, generate_wakepotential,
+ small_bunch):
wp = generate_wakepotential(plane="x")
wp.track(small_bunch)
-
+
# Test with default options
fig1 = wp.plot_last_wake('Wxdip')
assert fig1 is not None
-
+
# Test with plot_rho=False
fig2 = wp.plot_last_wake('Wxdip', plot_rho=False)
assert fig2 is not None
-
+
# Test with plot_dipole=True
fig3 = wp.plot_last_wake('Wxdip', plot_dipole=True)
assert fig3 is not None
-
+
# Test with plot_wake_function=False
fig4 = wp.plot_last_wake('Wxdip', plot_wake_function=False)
assert fig4 is not None
-
- @pytest.mark.parametrize("plane", [("x"), ("y"),("long")])
+
+ @pytest.mark.parametrize("plane", [("x"), ("y"), ("long")])
def test_get_gaussian_wakepotential(self, generate_wakepotential, plane):
wp = generate_wakepotential(plane=plane)
if plane == "x" or plane == "y":
comp = f"W{plane}dip"
else:
comp = "Wlong"
-
- tau0, W0, Wp, profile0, dipole0 = wp.get_gaussian_wakepotential(1e-12, wake_type=comp)
-
+
+ tau0, W0, Wp, profile0, dipole0 = wp.get_gaussian_wakepotential(
+ 1e-12, wake_type=comp)
+
assert tau0 is not None
assert W0 is not None
assert Wp is not None
assert profile0 is not None
assert dipole0 is not None
-
- @pytest.mark.parametrize("plane", [("x"), ("y"),("long")])
+
+ @pytest.mark.parametrize("plane", [("x"), ("y"), ("long")])
def test_plot_gaussian_wake(self, generate_wakepotential, plane):
wp = generate_wakepotential(plane=plane)
if plane == "x" or plane == "y":
@@ -198,9 +225,11 @@ class TestWakePotential:
comp = "Wlong"
fig = wp.plot_gaussian_wake(sigma=10e-12, wake_type=comp)
assert fig is not None
-
+
+
@pytest.fixture
def generate_lrrw(demo_ring, beam_uniform):
+
def generate(ring=demo_ring,
beam=beam_uniform,
length=demo_ring.L,
@@ -211,23 +240,26 @@ def generate_lrrw(demo_ring, beam_uniform):
x3=None,
y3=None,
average_beta=None):
- lrrw = LongRangeResistiveWall(ring=ring,
- beam=beam,
- length=length,
+ lrrw = LongRangeResistiveWall(ring=ring,
+ beam=beam,
+ length=length,
rho=rho,
radius=radius,
types=types,
- nt=nt,
+ nt=nt,
x3=x3,
y3=y3,
average_beta=average_beta)
return lrrw
+
return generate
+
class TestLongRangeResistiveWall:
-
+
# Test tracking w/ uniform beam
- @pytest.mark.parametrize("types, attr", [("Wxdip","xp"), ("Wydip","yp"), ("Wlong","delta")])
+ @pytest.mark.parametrize("types, attr", [("Wxdip", "xp"), ("Wydip", "yp"),
+ ("Wlong", "delta")])
def test_track(self, generate_lrrw, beam_uniform, types, attr):
lrrw = generate_lrrw(types=types)
for i, bunch in enumerate(beam_uniform.not_empty):
@@ -235,30 +267,37 @@ class TestLongRangeResistiveWall:
with np.errstate(over='ignore'):
lrrw.track(beam_uniform)
for i, bunch in enumerate(beam_uniform.not_empty):
- assert not np.allclose(getattr(self, "initial_attr_{i}"), bunch[attr])
-
+ assert not np.allclose(getattr(self, "initial_attr_{i}"),
+ bunch[attr])
+
# Test tracking w/ non-uniform beam
- @pytest.mark.parametrize("types, attr", [("Wxdip","xp"), ("Wydip","yp"), ("Wlong","delta")])
- def test_track_non_uniform(self, generate_lrrw, beam_non_uniform, types, attr):
+ @pytest.mark.parametrize("types, attr", [("Wxdip", "xp"), ("Wydip", "yp"),
+ ("Wlong", "delta")])
+ def test_track_non_uniform(self, generate_lrrw, beam_non_uniform, types,
+ attr):
lrrw = generate_lrrw(types=types, beam=beam_non_uniform)
for i, bunch in enumerate(beam_non_uniform.not_empty):
setattr(self, "initial_attr_{i}", bunch[attr].copy())
with np.errstate(over='ignore'):
lrrw.track(beam_non_uniform)
for i, bunch in enumerate(beam_non_uniform.not_empty):
- assert not np.allclose(getattr(self, "initial_attr_{i}"), bunch[attr])
-
+ assert not np.allclose(getattr(self, "initial_attr_{i}"),
+ bunch[attr])
+
# Test tracking w/ mpi beam
- @pytest.mark.parametrize("types, attr", [("Wxdip","xp"), ("Wydip","yp"), ("Wlong","delta")])
+ @pytest.mark.parametrize("types, attr", [("Wxdip", "xp"), ("Wydip", "yp"),
+ ("Wlong", "delta")])
def test_track_mpi(self, generate_lrrw, beam_1bunch_mpi, types, attr):
lrrw = generate_lrrw(types=types, beam=beam_1bunch_mpi)
bunch = beam_1bunch_mpi[beam_1bunch_mpi.mpi.bunch_num]
initial_attr = bunch[attr].copy()
with np.errstate(over='ignore'):
lrrw.track(beam_1bunch_mpi)
- lrrw.track(beam_1bunch_mpi) # two turns are needed to see the effect of the bunch on itself
+ lrrw.track(
+ beam_1bunch_mpi
+ ) # two turns are needed to see the effect of the bunch on itself
assert not np.allclose(initial_attr, bunch[attr])
-
+
# test that two turns are needed to see the effect of the bunch on itself
def test_kick_signle_bunch(self, generate_lrrw, beam_1bunch_mpi):
lrrw = generate_lrrw(beam=beam_1bunch_mpi)
@@ -293,7 +332,8 @@ class TestLongRangeResistiveWall:
assert not np.array_equal(lrrw.charge, initial_charge)
# Verify the behavior when the approximated wake functions are not valid
- def test_invalid_approximated_wake_functions(self, demo_ring, beam_uniform):
+ def test_invalid_approximated_wake_functions(self, demo_ring,
+ beam_uniform):
demo_ring.T1 = 1e-18 # Set T1 to a very small value to trigger the error
with pytest.raises(ValueError):
LongRangeResistiveWall(demo_ring, beam_uniform, 100, 1e-6, 6e-12)
@@ -306,25 +346,32 @@ class TestLongRangeResistiveWall:
initial_delta = bunch["delta"].copy()
initial_xp = bunch["xp"].copy()
initial_yp = bunch["yp"].copy()
-
+
# Act
lrrw.track_bunch(bunch, rank=0)
-
+
# Assert
if "Wlong" in lrrw.types:
- assert not np.array_equal(bunch["delta"], initial_delta), "Longitudinal kick not applied correctly"
+ assert not np.array_equal(
+ bunch["delta"],
+ initial_delta), "Longitudinal kick not applied correctly"
if "Wxdip" in lrrw.types:
- assert not np.array_equal(bunch["xp"], initial_xp), "Horizontal dipole kick not applied correctly"
+ assert not np.array_equal(
+ bunch["xp"],
+ initial_xp), "Horizontal dipole kick not applied correctly"
if "Wydip" in lrrw.types:
- assert not np.array_equal(bunch["yp"], initial_yp), "Vertical dipole kick not applied correctly"
+ assert not np.array_equal(
+ bunch["yp"],
+ initial_yp), "Vertical dipole kick not applied correctly"
def test_average_beta_local(self, generate_lrrw):
lrrw = generate_lrrw()
assert lrrw.norm_x == 1
assert lrrw.norm_y == 1
- def test_average_beta_mean(self, generate_lrrw, generate_ring_with_at_lattice):
- ring = generate_ring_with_at_lattice(tracking_loc = 5)
+ def test_average_beta_mean(self, generate_lrrw,
+ generate_ring_with_at_lattice):
+ ring = generate_ring_with_at_lattice(tracking_loc=5)
lrrw = generate_lrrw(ring=ring)
assert lrrw.norm_x != 1
- assert lrrw.norm_y != 1
\ No newline at end of file
+ assert lrrw.norm_y != 1
diff --git a/tests/unit/utilities/test_optics.py b/tests/unit/utilities/test_optics.py
index c5a0eea554e1a048420628c27cf89fd6f2fa5564..90e13aa5f03df0eec6a2a2256dd7dc7eff7b589a 100644
--- a/tests/unit/utilities/test_optics.py
+++ b/tests/unit/utilities/test_optics.py
@@ -1,22 +1,29 @@
import numpy as np
import pytest
+
pytestmark = pytest.mark.unit
+from pathlib import Path
+
import at
import matplotlib.pyplot as plt
-from pathlib import Path
+
from mbtrack2 import Optics, PhysicalModel
+
@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)
+ local_optics = Optics(local_beta=beta,
+ local_alpha=alpha,
+ local_dispersion=dispersion)
return local_optics
+
@pytest.fixture
def generate_at_optics():
+
def generate(lattice_file=None,
tracking_loc=0,
local_beta=None,
@@ -33,12 +40,15 @@ def generate_at_optics():
local_dispersion=local_dispersion,
**kwargs)
return optics
+
return generate
+
@pytest.fixture
def at_optics(generate_at_optics):
return generate_at_optics()
+
class TestOptics:
# Initialize Optics with a valid lattice file and verify attributes are set correctly
@@ -55,18 +65,30 @@ class TestOptics:
# Load a lattice using load_from_AT and verify optic functions are interpolated
def test_interpolation(self, at_optics):
lattice = at_optics.lattice
-
+
refpts = np.arange(0, len(lattice))
twiss0, tune, chrom, twiss = at.linopt(lattice,
refpts=refpts,
get_chrom=True)
randnum = np.random.randint(0, len(lattice))
pos = twiss.s_pos[randnum]
-
- np.testing.assert_allclose(at_optics.beta(pos), twiss.beta[randnum, :], rtol=1e-3, atol=1e-4)
- np.testing.assert_allclose(at_optics.alpha(pos), twiss.alpha[randnum, :], rtol=1e-3, atol=1e-4)
- np.testing.assert_allclose(at_optics.mu(pos), twiss.mu[randnum, :], rtol=1e-3, atol=1e-4)
- np.testing.assert_allclose(at_optics.dispersion(pos), twiss.dispersion[randnum, :], rtol=1e-3, atol=1e-4)
+
+ np.testing.assert_allclose(at_optics.beta(pos),
+ twiss.beta[randnum, :],
+ rtol=1e-3,
+ atol=1e-4)
+ np.testing.assert_allclose(at_optics.alpha(pos),
+ twiss.alpha[randnum, :],
+ rtol=1e-3,
+ atol=1e-4)
+ np.testing.assert_allclose(at_optics.mu(pos),
+ twiss.mu[randnum, :],
+ rtol=1e-3,
+ atol=1e-4)
+ np.testing.assert_allclose(at_optics.dispersion(pos),
+ twiss.dispersion[randnum, :],
+ rtol=1e-3,
+ atol=1e-4)
# Retrieve beta, alpha, gamma, dispersion, and mu functions at specified positions
def test_local_optic_functions(self, local_optics):
@@ -76,7 +98,7 @@ class TestOptics:
gamma = np.squeeze(local_optics.gamma(position))
dispersion = np.squeeze(local_optics.dispersion(position))
mu = np.squeeze(local_optics.mu(position))
-
+
np.testing.assert_allclose(beta, local_optics.local_beta)
np.testing.assert_allclose(alpha, local_optics.local_alpha)
np.testing.assert_allclose(gamma, local_optics.local_gamma)
@@ -92,15 +114,22 @@ class TestOptics:
def test_optics_tracking_loc(self, generate_at_optics, tracking_loc):
at_optics = generate_at_optics(tracking_loc=tracking_loc, n_points=1e4)
- np.testing.assert_allclose(at_optics.beta(tracking_loc), at_optics.local_beta)
- np.testing.assert_allclose(at_optics.alpha(tracking_loc), at_optics.local_alpha)
- np.testing.assert_allclose(at_optics.gamma(tracking_loc), at_optics.local_gamma, rtol=1e-3) # high interpol. error
- np.testing.assert_allclose(at_optics.dispersion(tracking_loc), at_optics.local_dispersion)
+ np.testing.assert_allclose(at_optics.beta(tracking_loc),
+ at_optics.local_beta)
+ np.testing.assert_allclose(at_optics.alpha(tracking_loc),
+ at_optics.local_alpha)
+ np.testing.assert_allclose(at_optics.gamma(tracking_loc),
+ at_optics.local_gamma,
+ rtol=1e-3) # high interpol. error
+ np.testing.assert_allclose(at_optics.dispersion(tracking_loc),
+ at_optics.local_dispersion)
+
class TestPhysicalModel:
-
+
@pytest.fixture
def generate_pm(self, ring_with_at_lattice):
+
def generate(ring=ring_with_at_lattice,
x_right=0.01,
y_top=0.02,
@@ -108,25 +137,25 @@ class TestPhysicalModel:
rho=1e-7,
x_left=None,
y_bottom=None,
- n_points=1e4
- ):
- pm = PhysicalModel(ring,
- x_right=x_right,
- y_top=y_top,
- shape=shape,
- rho=rho,
- x_left=x_left,
- y_bottom=y_bottom,
- n_points=n_points)
+ n_points=1e4):
+ pm = PhysicalModel(ring,
+ x_right=x_right,
+ y_top=y_top,
+ shape=shape,
+ rho=rho,
+ x_left=x_left,
+ y_bottom=y_bottom,
+ n_points=n_points)
return pm
+
return generate
# Initialize PhysicalModel with default symmetric aperture values and verify array shapes
def test_init_default_symmetric_aperture(self, generate_pm):
- n_points=1e3
+ n_points = 1e3
pm = generate_pm(n_points=n_points)
- assert pm.x_right.shape == (int(n_points),)
- assert pm.x_left.shape == (int(n_points),)
+ assert pm.x_right.shape == (int(n_points), )
+ assert pm.x_left.shape == (int(n_points), )
assert np.allclose(pm.x_left, -pm.x_right)
assert np.allclose(pm.y_bottom, -pm.y_top)
@@ -136,8 +165,8 @@ class TestPhysicalModel:
y_top_init = 0.02
shape_init = 'rect'
rho_init = 1e-7
- model = generate_pm(x_right=x_right_init,
- y_top=y_top_init,
+ model = generate_pm(x_right=x_right_init,
+ y_top=y_top_init,
shape=shape_init,
rho=rho_init)
x_right_modif = 0.015
@@ -146,16 +175,17 @@ class TestPhysicalModel:
rho_modif = 1e-6
start_position = 10
end_position = 20
- model.change_values(start_position,
- end_position,
- x_right=x_right_modif,
- y_top=y_top_modif,
+ model.change_values(start_position,
+ end_position,
+ x_right=x_right_modif,
+ y_top=y_top_modif,
shape=shape_modif,
rho=rho_modif)
- idx = (model.position > start_position) & (model.position < end_position)
+ idx = (model.position > start_position) & (model.position
+ < end_position)
idx2 = ((model.position[:-1] > start_position) &
(model.position[1:] < end_position))
-
+
# Assert modif
assert np.allclose(model.x_right[idx], x_right_modif)
assert np.allclose(model.x_left[idx], -x_right_modif)
@@ -170,7 +200,7 @@ class TestPhysicalModel:
assert np.allclose(model.y_bottom[~idx], -y_top_init)
assert np.all(model.shape[~idx2] == shape_init)
assert np.allclose(model.rho[~idx2], rho_init)
-
+
# Change values with sym=False and verify asymmetric updates
def test_change_values_asymmetric(self, generate_pm):
x_right_init = 0.01
@@ -187,7 +217,7 @@ class TestPhysicalModel:
y_top_init = 0.02
shape_init = 'rect'
rho_init = 1e-7
- model = generate_pm(x_right=x_right_init,
+ model = generate_pm(x_right=x_right_init,
y_top=y_top_init,
shape=shape_init,
rho=rho_init)
@@ -197,16 +227,17 @@ class TestPhysicalModel:
rho_modif = 1e-6
start_position = 10
end_position = 20
- model.taper(start_position,
- end_position,
- x_right_start=x_right_start,
+ model.taper(start_position,
+ end_position,
+ x_right_start=x_right_start,
x_right_end=x_right_end,
shape=shape_modif,
rho=rho_modif)
- idx = (model.position > start_position) & (model.position < end_position)
+ idx = (model.position > start_position) & (model.position
+ < end_position)
idx2 = ((model.position[:-1] > start_position) &
(model.position[1:] < end_position))
-
+
# Assert edge points
assert model.x_right[idx][0] == pytest.approx(x_right_start)
assert model.x_right[idx][-1] == pytest.approx(x_right_end)
@@ -226,14 +257,14 @@ class TestPhysicalModel:
assert np.allclose(model.y_bottom, -y_top_init)
assert np.all(model.shape[~idx2] == shape_init)
assert np.allclose(model.rho[~idx2], rho_init)
-
+
# Create tapered transition between two positions with different aperture values
def test_taper_nonsym_transition(self, generate_pm):
x_right_init = 0.01
y_top_init = 0.02
shape_init = 'rect'
rho_init = 1e-7
- model = generate_pm(x_right=x_right_init,
+ model = generate_pm(x_right=x_right_init,
y_top=y_top_init,
shape=shape_init,
rho=rho_init)
@@ -243,17 +274,18 @@ class TestPhysicalModel:
rho_modif = 1e-6
start_position = 10
end_position = 20
- model.taper(start_position,
- end_position,
- x_right_start=x_right_start,
+ model.taper(start_position,
+ end_position,
+ x_right_start=x_right_start,
x_right_end=x_right_end,
shape=shape_modif,
rho=rho_modif,
sym=False)
- idx = (model.position > start_position) & (model.position < end_position)
+ idx = (model.position > start_position) & (model.position
+ < end_position)
idx2 = ((model.position[:-1] > start_position) &
(model.position[1:] < end_position))
-
+
# Assert edge points
assert model.x_right[idx][0] == pytest.approx(x_right_start)
assert model.x_right[idx][-1] == pytest.approx(x_right_end)
@@ -271,9 +303,11 @@ class TestPhysicalModel:
assert np.allclose(model.rho[~idx2], rho_init)
# Calculate effective radius for resistive wall with default right/top settings
- def test_resistive_wall_effective_radius_default(self, generate_pm, ring_with_at_lattice):
+ def test_resistive_wall_effective_radius_default(self, generate_pm,
+ ring_with_at_lattice):
model = generate_pm()
- out = model.resistive_wall_effective_radius(ring_with_at_lattice.optics)
+ out = model.resistive_wall_effective_radius(
+ ring_with_at_lattice.optics)
for val in out:
assert val > 0
@@ -293,10 +327,12 @@ class TestPhysicalModel:
assert all(isinstance(x, float) for x in aperture)
# Handle zero resistivity sections in resistive_wall_effective_radius calculations
- def test_zero_resistivity_handling(self, generate_pm, ring_with_at_lattice):
+ def test_zero_resistivity_handling(self, generate_pm,
+ ring_with_at_lattice):
model = generate_pm()
model.change_values(10, 20, rho=0)
- out = model.resistive_wall_effective_radius(ring_with_at_lattice.optics)
+ out = model.resistive_wall_effective_radius(
+ ring_with_at_lattice.optics)
out[0] < ring_with_at_lattice.L
for val in out:
assert val > 0
@@ -319,4 +355,7 @@ class TestPhysicalModel:
with pytest.raises(ValueError):
model.taper(-10, 10, x_right_start=0.01, x_right_end=0.03)
with pytest.raises(ValueError):
- model.taper(0, model.position[-1] + 10, x_right_start=0.01, x_right_end=0.03)
\ No newline at end of file
+ model.taper(0,
+ model.position[-1] + 10,
+ x_right_start=0.01,
+ x_right_end=0.03)
diff --git a/tests/utility_test_functions.py b/tests/utility_test_functions.py
index 1be60451aadb2ced45bde82ceb43f940b76b16d0..0ab7377fbf2fd28d59670a014e6002836db2dce8 100644
--- a/tests/utility_test_functions.py
+++ b/tests/utility_test_functions.py
@@ -1,11 +1,13 @@
import numpy as np
-from mbtrack2 import Bunch, Beam
-def assert_attr_changed(element,
- bunch,
+from mbtrack2 import Beam, Bunch
+
+
+def assert_attr_changed(element,
+ bunch,
attrs_changed=["xp", "yp", "delta"],
change=True):
-
+
if isinstance(bunch, Bunch):
assert_attr_changed_bunch(element, bunch, attrs_changed, change=change)
elif isinstance(bunch, Beam):
@@ -13,51 +15,71 @@ def assert_attr_changed(element,
else:
raise TypeError
-def assert_attr_changed_bunch(element,
- bunch,
+
+def assert_attr_changed_bunch(element,
+ bunch,
attrs_changed=["xp", "yp", "delta"],
change=True):
-
- attrs = ["x","xp","y","yp","tau","delta"]
+
+ attrs = ["x", "xp", "y", "yp", "tau", "delta"]
attrs_unchanged = [attr for attr in attrs if attr not in attrs_changed]
-
- initial_values_changed = {attr: bunch[attr].copy() for attr in attrs_changed}
-
- initial_values_unchanged = {attr: bunch[attr].copy() for attr in attrs_unchanged}
-
+
+ initial_values_changed = {
+ attr: bunch[attr].copy()
+ for attr in attrs_changed
+ }
+
+ initial_values_unchanged = {
+ attr: bunch[attr].copy()
+ for attr in attrs_unchanged
+ }
+
element.track(bunch)
for attr in attrs_changed:
if change:
- assert not np.array_equal(initial_values_changed[attr], bunch[attr]), f"{attr}"
+ assert not np.array_equal(initial_values_changed[attr],
+ bunch[attr]), f"{attr}"
else:
- assert np.array_equal(initial_values_changed[attr], bunch[attr]), f"{attr}"
-
+ assert np.array_equal(initial_values_changed[attr],
+ bunch[attr]), f"{attr}"
+
for attr in attrs_unchanged:
- assert np.array_equal(initial_values_unchanged[attr], bunch[attr]), f"{attr}"
-
-def assert_attr_changed_beam(element,
+ assert np.array_equal(initial_values_unchanged[attr],
+ bunch[attr]), f"{attr}"
+
+
+def assert_attr_changed_beam(element,
beam,
attrs_changed=["xp", "yp", "delta"],
change=True):
-
- attrs = ["x","xp","y","yp","tau","delta"]
+
+ attrs = ["x", "xp", "y", "yp", "tau", "delta"]
attrs_unchanged = [attr for attr in attrs if attr not in attrs_changed]
-
- initial_values_changed_b = [{attr: bunch[attr].copy() for attr in attrs_changed} for bunch in beam]
-
- initial_values_unchanged_b = [{attr: bunch[attr].copy() for attr in attrs_unchanged} for bunch in beam]
-
+
+ initial_values_changed_b = [{
+ attr: bunch[attr].copy()
+ for attr in attrs_changed
+ } for bunch in beam]
+
+ initial_values_unchanged_b = [{
+ attr: bunch[attr].copy()
+ for attr in attrs_unchanged
+ } for bunch in beam]
+
element.track(beam)
-
+
for i, bunch in enumerate(beam):
initial_values_changed = initial_values_changed_b[i]
initial_values_unchanged = initial_values_unchanged_b[i]
for attr in attrs_changed:
if change and (bunch.charge != 0):
- assert not np.array_equal(initial_values_changed[attr], bunch[attr]), f"{attr}"
+ assert not np.array_equal(initial_values_changed[attr],
+ bunch[attr]), f"{attr}"
else:
- assert np.array_equal(initial_values_changed[attr], bunch[attr]), f"{attr}"
-
+ assert np.array_equal(initial_values_changed[attr],
+ bunch[attr]), f"{attr}"
+
for attr in attrs_unchanged:
- assert np.array_equal(initial_values_unchanged[attr], bunch[attr]), f"{attr}"
\ No newline at end of file
+ assert np.array_equal(initial_values_unchanged[attr],
+ bunch[attr]), f"{attr}"
diff --git a/tests/utility_test_growth_rate.py b/tests/utility_test_growth_rate.py
index cd3fee16b676072b6417e6e8f986d1c799e04065..8c37ac1a791eed98916c2ea1fd066f98a2e41246 100644
--- a/tests/utility_test_growth_rate.py
+++ b/tests/utility_test_growth_rate.py
@@ -1,50 +1,57 @@
-import numpy as np
-import matplotlib.pyplot as plt
import h5py as hp
+import matplotlib.pyplot as plt
+import numpy as np
+
def r2(y, y_fit):
# residual sum of squares
- ss_res = np.sum((y - y_fit) ** 2)
+ ss_res = np.sum((y - y_fit)**2)
# total sum of squares
- ss_tot = np.sum((y - np.mean(y)) ** 2)
+ ss_tot = np.sum((y - np.mean(y))**2)
# r-squared
- r2 = 1 - (ss_res / ss_tot)
+ r2 = 1 - (ss_res/ss_tot)
return r2
+
def get_growth(ring, start, end, bunch_num, g, plane="y", plot=False):
- plane_dict = {"x":0, "y":1, "long":2}
- time = np.array(g["Beam"]["time"])*ring.T0
- data = np.array(g["Beam"]["cs_invariant"][plane_dict[plane],bunch_num,:])
- idx = (time > ring.T0*start) & (time < ring.T0*end)
- time_fit = time[idx] - start*ring.T0
+ plane_dict = {"x": 0, "y": 1, "long": 2}
+ time = np.array(g["Beam"]["time"]) * ring.T0
+ data = np.array(g["Beam"]["cs_invariant"][plane_dict[plane], bunch_num, :])
+ idx = (time > ring.T0 * start) & (time < ring.T0 * end)
+ time_fit = time[idx] - start * ring.T0
data_fit = data[idx]
- ploy = np.polynomial.polynomial.Polynomial.fit(time_fit, np.log(data_fit), 1, w=np.sqrt(data_fit))
-
+ ploy = np.polynomial.polynomial.Polynomial.fit(time_fit,
+ np.log(data_fit),
+ 1,
+ w=np.sqrt(data_fit))
+
if plot:
- plt.plot(time/ring.T0, np.log(data), '-')
+ plt.plot(time / ring.T0, np.log(data), '-')
arr = ploy.convert().coef
- fit_vals = [np.exp(arr[0]) * np.exp(arr[1] *x) for x in time_fit]
-
- return (arr[0], arr[1]/2, r2(data_fit, fit_vals))
+ fit_vals = [np.exp(arr[0]) * np.exp(arr[1] * x) for x in time_fit]
+
+ return (arr[0], arr[1] / 2, r2(data_fit, fit_vals))
+
def get_growth_beam(ring, path, start, end, plane="y", plot=False):
- g = hp.File(path,"r")
+ g = hp.File(path, "r")
gr_cs = np.zeros(ring.h)
cste = np.zeros(ring.h)
r2_cs = np.zeros(ring.h)
for i in range(ring.h):
- cste[i], gr_cs[i], r2_cs[i] = get_growth(ring, start, end, i, g, plane, plot)
-
+ cste[i], gr_cs[i], r2_cs[i] = get_growth(ring, start, end, i, g, plane,
+ plot)
+
time = np.array(g["Beam"]["time"])
idx = (time > start) & (time < end)
time_fit = time[idx]
- fit = np.mean(cste) + np.mean(gr_cs)*2*(time_fit-start)*ring.T0
-
+ fit = np.mean(cste) + np.mean(gr_cs) * 2 * (time_fit-start) * ring.T0
+
print(f"growth rate = {np.mean(gr_cs)} +- {np.std(gr_cs)} s-1")
print(f"r2 = {np.mean(r2_cs)} +- {np.std(r2_cs)}")
if plot:
@@ -57,4 +64,4 @@ def get_growth_beam(ring, path, start, end, plane="y", plot=False):
plt.xlabel("Turn number")
plt.ylabel("log(CS_invariant)")
g.close()
- return (np.mean(gr_cs), np.std(gr_cs))
\ No newline at end of file
+ return (np.mean(gr_cs), np.std(gr_cs))