Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions orbit_tools/bunch.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,14 @@
from .cov import normalization_matrix


def set_particle_macrosizes(bunch: Bunch, macrosizes: Iterable) -> Bunch:
bunch.addPartAttr("macrosize") # sets macrosize=0 for all particles
attribute_array_index = 0
for index in range(bunch.getSize()):
bunch.partAttrValue("macrosize", index, attribute_array_index, macrosizes[index])
return bunch


def get_part_coords(bunch: Bunch, index: int) -> list[float]:
x = bunch.x(index)
y = bunch.y(index)
Expand Down
9 changes: 9 additions & 0 deletions tests/test_bunch.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,3 +126,12 @@ def sample():
def test_get_bunch_twiss_containers():
bunch = make_gaussian_bunch()
twiss_x, twiss_y, twiss_z = ot.bunch.get_bunch_twiss_containers(bunch)


def test_set_particle_macrosizes():
bunch = make_gaussian_bunch()
macrosizes = list(range(bunch.getSize()))
bunch = ot.bunch.set_particle_macrosizes(bunch, macrosizes)
for i, macrosize in enumerate(macrosizes):
assert macrosize == bunch.partAttrValue("macrosize", i, 0)

Loading