Skip to content

Commit 5bf94c9

Browse files
giadarolgiadarol
committed
Merge branch 'feature/aperture_interpolation' into main
2 parents a438027 + 3f60d7c commit 5bf94c9

22 files changed

+1680
-133
lines changed

‎examples/apertures/003_polygon.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import numpy as np
2+
3+
import xobjects as xo
4+
import xtrack as xt
5+
6+
context = xo.ContextPyopencl()
7+
np2ctx = context.nparray_to_context_array
8+
ctx2np = context.nparray_from_context_array
9+
10+
11+
x_vertices=np.array([1.5, 0.2, -1, -1, 1])*1e-2
12+
y_vertices=np.array([1.3, 0.5, 1, -1, -1])*1e-2
13+
14+
aper = xt.LimitPolygon(
15+
_context=context,
16+
x_vertices=np2ctx(x_vertices),
17+
y_vertices=np2ctx(y_vertices))
18+
19+
# Try some particles inside
20+
parttest = xt.Particles(
21+
_context=context,
22+
p0c=6500e9,
23+
x=x_vertices*0.99,
24+
y=y_vertices*0.99)
25+
aper.track(parttest)
26+
assert np.allclose(ctx2np(parttest.state), 1)
27+
28+
# Try some particles outside
29+
parttest = xt.Particles(
30+
_context=context,
31+
p0c=6500e9,
32+
x=x_vertices*1.01,
33+
y=y_vertices*1.01)
34+
aper.track(parttest)
35+
assert np.allclose(ctx2np(parttest.state), 0)
36+
37+
part_gen_range = 0.02
38+
n_part=10000
39+
particles = xt.Particles(
40+
_context=context,
41+
p0c=6500e9,
42+
x=np.random.uniform(-part_gen_range, part_gen_range, n_part),
43+
px = np.zeros(n_part),
44+
y=np.random.uniform(-part_gen_range, part_gen_range, n_part),
45+
py = np.zeros(n_part),
46+
sigma = np.zeros(n_part),
47+
delta = np.zeros(n_part))
48+
49+
aper.track(particles)
50+
51+
part_id = context.nparray_from_context_array(particles.particle_id)
52+
part_state = context.nparray_from_context_array(particles.state)
53+
part_x = context.nparray_from_context_array(particles.x)
54+
part_y = context.nparray_from_context_array(particles.y)
55+
56+
import matplotlib.pyplot as plt
57+
plt.close('all')
58+
fig1 = plt.figure(1)
59+
ax = fig1.add_subplot(111)
60+
plt.plot(part_x, part_y, '.', color='red')
61+
plt.plot(part_x[part_state>0], part_y[part_state>0], '.', color='green')
62+
ax.plot(aper.x_closed, aper.y_closed, linewidth=3, color='black')
63+
64+
plt.show()

‎examples/collimation/001_loss_location_refinement.py‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
import logging
2+
3+
import numpy as np
4+
import matplotlib.pyplot as plt
5+
6+
import xline as xl
7+
import xtrack as xt
8+
import xobjects as xo
9+
import xpart as xp
10+
from scipy.spatial import ConvexHull
11+
12+
13+
plt.close('all')
14+
15+
n_part=10000
16+
ctx = xo.context_default
17+
buf = ctx.new_buffer()
18+
19+
logger = logging.getLogger('xtrack')
20+
logger.setLevel(logging.DEBUG)
21+
22+
# Define aper_0
23+
#aper_0 = xt.LimitRect(_buffer=buf, min_y=-1e-2, max_y=1e-2,
24+
# min_x=-2e-2, max_x=2e-2)
25+
aper_0 = xt.LimitEllipse(_buffer=buf, a=2e-2, b=1e-2)
26+
shift_aper_0 = (1e-2, 0.5e-2)
27+
rot_deg_aper_0 = 10.
28+
29+
# Define aper_1
30+
#aper_1 = xt.LimitEllipse(_buffer=buf, a=1e-2, b=2e-2)
31+
aper_1 = xt.LimitRect(_buffer=buf, min_x=-1e-2, max_x=1e-2,
32+
min_y=-2e-2, max_y=2e-2)
33+
shift_aper_1 = (-5e-3, 1e-2)
34+
rot_deg_aper_1 = 10.
35+
36+
37+
# aper_0_sandwitch
38+
trk_aper_0 = xt.Tracker(_buffer=buf, sequence=xl.Line(
39+
elements=[xt.XYShift(_buffer=buf, dx=shift_aper_0[0], dy=shift_aper_0[1]),
40+
xt.SRotation(_buffer=buf, angle=rot_deg_aper_0),
41+
aper_0,
42+
xt.Multipole(_buffer=buf, knl=[0.001]),
43+
xt.SRotation(_buffer=buf, angle=-rot_deg_aper_0),
44+
xt.XYShift(_buffer=buf, dx=-shift_aper_0[0], dy=-shift_aper_0[1])]))
45+
46+
# aper_1_sandwitch
47+
trk_aper_1 = xt.Tracker(_buffer=buf, sequence=xl.Line(
48+
elements=[xt.XYShift(_buffer=buf, dx=shift_aper_1[0], dy=shift_aper_1[1]),
49+
xt.SRotation(_buffer=buf, angle=rot_deg_aper_1),
50+
aper_1,
51+
xt.Multipole(_buffer=buf, knl=[0.001]),
52+
xt.SRotation(_buffer=buf, angle=-rot_deg_aper_1),
53+
xt.XYShift(_buffer=buf, dx=-shift_aper_1[0], dy=-shift_aper_1[1])]))
54+
55+
# Build example line
56+
tracker = xt.Tracker(_buffer=buf, sequence=xl.Line(
57+
elements = ((xt.Drift(_buffer=buf, length=0.5),)
58+
+ trk_aper_0.line.elements
59+
+ (xt.Drift(_buffer=buf, length=1),
60+
xt.Multipole(_buffer=buf, knl=[1e-3]),
61+
xt.Drift(_buffer=buf, length=1),
62+
xt.Cavity(_buffer=buf, voltage=3e6, frequency=400e6),
63+
xt.Drift(_buffer=buf, length=1.),)
64+
+ trk_aper_1.line.elements)))
65+
num_elements = len(tracker.line.elements)
66+
67+
# Test on full line
68+
particles = xt.Particles(_context=ctx,
69+
px=np.random.uniform(-0.01, 0.01, 10000),
70+
py=np.random.uniform(-0.01, 0.01, 10000))
71+
72+
tracker.track(particles)
73+
74+
75+
loss_loc_refinement = xt.LossLocationRefinement(tracker,
76+
n_theta = 360,
77+
r_max = 0.5, # m
78+
dr = 50e-6,
79+
ds = 0.1,
80+
save_refine_trackers=True)
81+
82+
import time
83+
t0 = time.time()
84+
85+
loss_loc_refinement.refine_loss_location(particles)
86+
87+
t1 = time.time()
88+
print(f'Took\t{(t1-t0)*1e3:.2f} ms')
89+
90+
# Visualize apertures
91+
interp_tracker = loss_loc_refinement.refine_trackers[
92+
loss_loc_refinement.i_apertures[1]]
93+
s0 = interp_tracker.s0
94+
s1 = interp_tracker.s1
95+
polygon_0 = interp_tracker.line.elements[0]
96+
polygon_1 = interp_tracker.line.elements[-1]
97+
for ii, (trkr, poly) in enumerate(
98+
zip([trk_aper_0, trk_aper_1],
99+
[polygon_0, polygon_1])):
100+
part_gen_range = 0.05
101+
pp = xt.Particles(
102+
p0c=6500e9,
103+
x=np.random.uniform(-part_gen_range, part_gen_range, n_part),
104+
y=np.random.uniform(-part_gen_range, part_gen_range, n_part))
105+
x0 = pp.x.copy()
106+
y0 = pp.y.copy()
107+
108+
trkr.track(pp)
109+
ids = pp.particle_id
110+
111+
112+
fig = plt.figure(ii+1)
113+
ax = fig.add_subplot(111)
114+
plt.plot(x0, y0, '.', color='red')
115+
plt.axis('equal')
116+
plt.grid(linestyle=':')
117+
plt.plot(x0[ids][pp.state>0], y0[ids][pp.state>0], '.', color='green')
118+
plt.plot(poly.x_closed, poly.y_closed, '-k', linewidth=1)
119+
120+
plt.figure()
121+
ax = plt.axes(projection='3d')
122+
ax.plot3D(
123+
polygon_0.x_closed,
124+
polygon_0.y_closed,
125+
s0+polygon_0.x_closed*0,
126+
color='k', linewidth=3)
127+
ax.plot3D(
128+
polygon_1.x_closed,
129+
polygon_1.y_closed,
130+
s1+polygon_1.x_closed*0,
131+
color='k', linewidth=3)
132+
s_check = []
133+
for ee, ss in zip(interp_tracker.line.elements,
134+
interp_tracker.line.element_s_locations):
135+
if ee.__class__ is xt.LimitPolygon:
136+
ax.plot3D(
137+
ee.x_closed,
138+
ee.y_closed,
139+
s0+ss+0*ee.x_closed,
140+
alpha=0.9,
141+
)
142+
s_check.append(ss)
143+
mask = particles.at_element == loss_loc_refinement.i_apertures[1]
144+
ax.plot3D(particles.x[mask], particles.y[mask], particles.s[mask],
145+
'.r', markersize=.5)
146+
ax.view_init(65, 62); plt.draw()
147+
plt.show()

‎examples/collimation/002_loss_refinement_check_cone.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
import logging
2+
3+
import numpy as np
4+
import matplotlib.pyplot as plt
5+
6+
import xline as xl
7+
import xtrack as xt
8+
import xobjects as xo
9+
import xpart as xp
10+
from scipy.spatial import ConvexHull
11+
12+
13+
plt.close('all')
14+
15+
n_part=10000
16+
shift_x = 0.3e-2
17+
shift_y = 0.5e-2
18+
19+
ctx = xo.context_default
20+
buf = ctx.new_buffer()
21+
22+
logger = logging.getLogger('xtrack')
23+
logger.setLevel(logging.DEBUG)
24+
25+
# Define aper_0
26+
aper_0 = xt.LimitEllipse(_buffer=buf, a=2e-2, b=2e-2)
27+
shift_aper_0 = (shift_x, shift_y)
28+
rot_deg_aper_0 = 10.
29+
30+
# Define aper_1
31+
aper_1 = xt.LimitEllipse(_buffer=buf, a=1e-2, b=1e-2)
32+
shift_aper_1 = (shift_x, shift_y)
33+
rot_deg_aper_1 = 10.
34+
35+
# aper_0_sandwitch
36+
trk_aper_0 = xt.Tracker(_buffer=buf, sequence=xl.Line(
37+
elements=[xt.XYShift(_buffer=buf, dx=shift_aper_0[0], dy=shift_aper_0[1]),
38+
xt.SRotation(_buffer=buf, angle=rot_deg_aper_0),
39+
aper_0,
40+
xt.Multipole(_buffer=buf, knl=[0.00]),
41+
xt.SRotation(_buffer=buf, angle=-rot_deg_aper_0),
42+
xt.XYShift(_buffer=buf, dx=-shift_aper_0[0], dy=-shift_aper_0[1])]))
43+
44+
# aper_1_sandwitch
45+
trk_aper_1 = xt.Tracker(_buffer=buf, sequence=xl.Line(
46+
elements=[xt.XYShift(_buffer=buf, dx=shift_aper_1[0], dy=shift_aper_1[1]),
47+
xt.SRotation(_buffer=buf, angle=rot_deg_aper_1),
48+
aper_1,
49+
xt.Multipole(_buffer=buf, knl=[0.00]),
50+
xt.SRotation(_buffer=buf, angle=-rot_deg_aper_1),
51+
xt.XYShift(_buffer=buf, dx=-shift_aper_1[0], dy=-shift_aper_1[1])]))
52+
53+
# Build example line
54+
tracker = xt.Tracker(_buffer=buf, sequence=xl.Line(
55+
elements = ((xt.Drift(_buffer=buf, length=0.5),)
56+
+ trk_aper_0.line.elements
57+
+ (xt.Drift(_buffer=buf, length=1),
58+
xt.Multipole(_buffer=buf, knl=[0.]),
59+
xt.Drift(_buffer=buf, length=1),
60+
xt.Cavity(_buffer=buf, voltage=3e6, frequency=400e6),
61+
xt.Drift(_buffer=buf, length=1.),)
62+
+ trk_aper_1.line.elements)))
63+
num_elements = len(tracker.line.elements)
64+
65+
# Test on full line
66+
r = np.linspace(0, 0.018, n_part)
67+
theta = np.linspace(0, 8*np.pi, n_part)
68+
particles = xt.Particles(_context=ctx,
69+
p0c=6500e9,
70+
x=r*np.cos(theta)+shift_x,
71+
y=r*np.sin(theta)+shift_y)
72+
73+
tracker.track(particles)
74+
75+
76+
loss_loc_refinement = xt.LossLocationRefinement(tracker,
77+
n_theta = 360,
78+
r_max = 0.5, # m
79+
dr = 50e-6,
80+
ds = 0.1,
81+
save_refine_trackers=True)
82+
83+
import time
84+
t0 = time.time()
85+
86+
loss_loc_refinement.refine_loss_location(particles)
87+
88+
t1 = time.time()
89+
print(f'Took\t{(t1-t0)*1e3:.2f} ms')
90+
91+
92+
# Automatic checks
93+
mask_lost = particles.state == 0
94+
r_calc = np.sqrt((particles.x-shift_x)**2 + (particles.y-shift_y)**2)
95+
assert np.all(r_calc[~mask_lost]<1e-2)
96+
assert np.all(r_calc[mask_lost]>1e-2)
97+
i_aper_1 = tracker.line.elements.index(aper_1)
98+
assert np.all(particles.at_element[mask_lost]==i_aper_1)
99+
assert np.all(particles.at_element[~mask_lost]==0)
100+
s0 = tracker.line.element_s_locations[tracker.line.elements.index(aper_0)]
101+
s1 = tracker.line.element_s_locations[tracker.line.elements.index(aper_1)]
102+
r0 = np.sqrt(aper_0.a_squ)
103+
r1 = np.sqrt(aper_1.a_squ)
104+
s_expected = s0 + (r_calc-r0)/(r1 - r0)*(s1 - s0)
105+
# TODO This threshold is a bit large
106+
assert np.allclose(particles.s[mask_lost], s_expected[mask_lost], atol=0.11)
107+
108+
# Visualize apertures
109+
interp_tracker = loss_loc_refinement.refine_trackers[
110+
loss_loc_refinement.i_apertures[1]]
111+
s0 = interp_tracker.s0
112+
s1 = interp_tracker.s1
113+
polygon_0 = interp_tracker.line.elements[0]
114+
polygon_1 = interp_tracker.line.elements[-1]
115+
for ii, (trkr, poly) in enumerate(
116+
zip([trk_aper_0, trk_aper_1],
117+
[polygon_0, polygon_1])):
118+
part_gen_range = 0.05
119+
pp = xt.Particles(
120+
p0c=6500e9,
121+
x=np.random.uniform(-part_gen_range, part_gen_range, n_part),
122+
y=np.random.uniform(-part_gen_range, part_gen_range, n_part))
123+
x0 = pp.x.copy()
124+
y0 = pp.y.copy()
125+
126+
trkr.track(pp)
127+
ids = pp.particle_id
128+
129+
130+
fig = plt.figure(ii+1)
131+
ax = fig.add_subplot(111)
132+
plt.plot(x0, y0, '.', color='red')
133+
plt.axis('equal')
134+
plt.grid(linestyle=':')
135+
plt.plot(x0[ids][pp.state>0], y0[ids][pp.state>0], '.', color='green')
136+
plt.plot(poly.x_closed, poly.y_closed, '-k', linewidth=1)
137+
138+
plt.figure()
139+
ax = plt.axes(projection='3d')
140+
ax.plot3D(
141+
polygon_0.x_closed,
142+
polygon_0.y_closed,
143+
s0+polygon_0.x_closed*0,
144+
color='k', linewidth=3)
145+
ax.plot3D(
146+
polygon_1.x_closed,
147+
polygon_1.y_closed,
148+
s1+polygon_1.x_closed*0,
149+
color='k', linewidth=3)
150+
s_check = []
151+
r_check = []
152+
for ee, ss in zip(interp_tracker.line.elements,
153+
interp_tracker.line.element_s_locations):
154+
if ee.__class__ is xt.LimitPolygon:
155+
ax.plot3D(
156+
ee.x_closed,
157+
ee.y_closed,
158+
s0+ss+0*ee.x_closed,
159+
alpha=0.9,
160+
)
161+
s_check.append(ss+s0)
162+
r_vertices = np.sqrt(
163+
(ee.x_vertices-shift_x)**2 + (ee.y_vertices-shift_y)**2)
164+
r_check.append(np.mean(r_vertices))
165+
mask = particles.at_element == loss_loc_refinement.i_apertures[1]
166+
ax.plot3D(particles.x[mask], particles.y[mask], particles.s[mask],
167+
'.r', markersize=.5)
168+
ax.view_init(65, 62); plt.draw()
169+
plt.show()

0 commit comments

Comments
 (0)