Skip to content

Commit 92f5e4c

Browse files
authored
Merge pull request #142 from xsuite/PencilPatch
Added kwargs to xp.pencil calls (such that e.g. method='4d' can be pa…
2 parents db216a6 + 8dbf3c3 commit 92f5e4c

File tree

1 file changed

+10
-6
lines changed

1 file changed

+10
-6
lines changed

xcoll/initial_distribution.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# copyright ############################### #
22
# This file is part of the Xcoll package. #
3-
# Copyright (c) CERN, 2024. #
3+
# Copyright (c) CERN, 2025. #
44
# ######################################### #
55

66
import numpy as np
@@ -88,11 +88,14 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
8888
if side == '+-':
8989
num_plus = int(num_particles/2)
9090
num_min = int(num_particles - num_plus)
91-
coords_plus = _generate_4D_pencil_one_jaw(line, name, num_plus, plane, '+', impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
92-
coords_min = _generate_4D_pencil_one_jaw(line, name, num_min, plane, '-', impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
91+
coords_plus = _generate_4D_pencil_one_jaw(line, name, num_plus, plane, '+', impact_parameter, dr_sigmas,
92+
at_element, is_converging, tw_at_s, **kwargs)
93+
coords_min = _generate_4D_pencil_one_jaw(line, name, num_min, plane, '-', impact_parameter, dr_sigmas,
94+
at_element, is_converging, tw_at_s, **kwargs)
9395
coords = [ [*c_plus, *c_min] for c_plus, c_min in zip(coords_plus, coords_min)]
9496
else:
95-
coords = _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter, dr_sigmas, at_element, is_converging, tw_at_s)
97+
coords = _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter,
98+
dr_sigmas, at_element, is_converging, tw_at_s, **kwargs)
9699
pencil = coords[0]
97100
p_pencil = coords[1]
98101
transverse_norm = coords[2]
@@ -183,7 +186,8 @@ def generate_delta_from_dispersion(line, at_element, *, plane, position_mm, nemi
183186

184187

185188
def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter,
186-
dr_sigmas, at_element, is_converging, tw_at_s=None):
189+
dr_sigmas, at_element, is_converging, tw_at_s=None,
190+
_capacity=None, **kwargs):
187191
coll = line[name]
188192

189193
if side == '+':
@@ -209,7 +213,7 @@ def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_p
209213
pencil, p_pencil = xp.generate_2D_pencil_with_absolute_cut(
210214
num_particles, plane=plane, absolute_cut=pencil_pos, line=line,
211215
dr_sigmas=dr_sigmas, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
212-
at_element=at_element, side=side, twiss=tw_at_s
216+
at_element=at_element, side=side, twiss=tw_at_s, **kwargs
213217
)
214218

215219
# Other plane: generate gaussian distribution in normalized coordinates

0 commit comments

Comments
 (0)