Something went wrong on our end
-
Alexis GAMELIN authoredCode modifications: * [Fix] adts in TransverseMapSector * ComplexData.name_and_coefficients_table is now a staticmethod. * Extend WakeField to "mixed" componenets (xydip, yxdip, ...). * Small fix/improvements for rf.py (CavityResonator for n_bin=1) * [Fix] WakePotential.check_sampling method. * [Fix] Sweep.plot method. * Add Element.track_bunch_if_non_empty decorator. * Modify Bunch.binning to add a return_full_length optional parameter. * Modify Mpi.share_distributions for n_bin=1. * Modify TransverseSpaceCharge to work with track_alive = True. * Modify WakePotential for n_bin=1 and use track_bunch_if_non_empty decorator. * [Fix] IntrabeamScattering for nonuniform fills with Beam.mpi_switch == False * Modify PhysicalModel to return errors if needed. * Update BeamIonElement Project management: * Update pyproject.toml and poetry.lock with pytest-mock * Add assert_attr_changed test function.
Alexis GAMELIN authoredCode modifications: * [Fix] adts in TransverseMapSector * ComplexData.name_and_coefficients_table is now a staticmethod. * Extend WakeField to "mixed" componenets (xydip, yxdip, ...). * Small fix/improvements for rf.py (CavityResonator for n_bin=1) * [Fix] WakePotential.check_sampling method. * [Fix] Sweep.plot method. * Add Element.track_bunch_if_non_empty decorator. * Modify Bunch.binning to add a return_full_length optional parameter. * Modify Mpi.share_distributions for n_bin=1. * Modify TransverseSpaceCharge to work with track_alive = True. * Modify WakePotential for n_bin=1 and use track_bunch_if_non_empty decorator. * [Fix] IntrabeamScattering for nonuniform fills with Beam.mpi_switch == False * Modify PhysicalModel to return errors if needed. * Update BeamIonElement Project management: * Update pyproject.toml and poetry.lock with pytest-mock * Add assert_attr_changed test function.
utility_test_functions.py 2.49 KiB
import numpy as np
from mbtrack2 import Bunch, Beam
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):
assert_attr_changed_beam(element, bunch, attrs_changed, change=change)
else:
raise TypeError
def assert_attr_changed_bunch(element,
bunch,
attrs_changed=["xp", "yp", "delta"],
change=True):
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}
element.track(bunch)
for attr in attrs_changed:
if change:
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}"
for attr in attrs_unchanged:
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_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]
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}"
else:
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}"