xcoll 0.3.6__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (56) hide show
  1. xcoll/__init__.py +12 -4
  2. xcoll/beam_elements/__init__.py +7 -5
  3. xcoll/beam_elements/absorber.py +41 -7
  4. xcoll/beam_elements/base.py +1161 -244
  5. xcoll/beam_elements/collimators_src/black_absorber.h +118 -0
  6. xcoll/beam_elements/collimators_src/black_crystal.h +111 -0
  7. xcoll/beam_elements/collimators_src/everest_block.h +40 -28
  8. xcoll/beam_elements/collimators_src/everest_collimator.h +129 -50
  9. xcoll/beam_elements/collimators_src/everest_crystal.h +217 -73
  10. xcoll/beam_elements/everest.py +60 -113
  11. xcoll/colldb.py +250 -750
  12. xcoll/general.py +2 -2
  13. xcoll/headers/checks.h +1 -1
  14. xcoll/headers/particle_states.h +2 -2
  15. xcoll/initial_distribution.py +195 -0
  16. xcoll/install.py +177 -0
  17. xcoll/interaction_record/__init__.py +1 -0
  18. xcoll/interaction_record/interaction_record.py +252 -0
  19. xcoll/interaction_record/interaction_record_src/interaction_record.h +98 -0
  20. xcoll/{impacts → interaction_record}/interaction_types.py +11 -4
  21. xcoll/line_tools.py +83 -0
  22. xcoll/lossmap.py +209 -0
  23. xcoll/manager.py +2 -937
  24. xcoll/rf_sweep.py +1 -1
  25. xcoll/scattering_routines/everest/amorphous.h +239 -0
  26. xcoll/scattering_routines/everest/channeling.h +245 -0
  27. xcoll/scattering_routines/everest/crystal_parameters.h +137 -0
  28. xcoll/scattering_routines/everest/everest.h +8 -30
  29. xcoll/scattering_routines/everest/everest.py +13 -10
  30. xcoll/scattering_routines/everest/jaw.h +27 -197
  31. xcoll/scattering_routines/everest/materials.py +2 -0
  32. xcoll/scattering_routines/everest/multiple_coulomb_scattering.h +31 -10
  33. xcoll/scattering_routines/everest/nuclear_interaction.h +86 -0
  34. xcoll/scattering_routines/geometry/__init__.py +6 -0
  35. xcoll/scattering_routines/geometry/collimator_geometry.h +219 -0
  36. xcoll/scattering_routines/geometry/crystal_geometry.h +150 -0
  37. xcoll/scattering_routines/geometry/geometry.py +26 -0
  38. xcoll/scattering_routines/geometry/get_s.h +92 -0
  39. xcoll/scattering_routines/geometry/methods.h +111 -0
  40. xcoll/scattering_routines/geometry/objects.h +154 -0
  41. xcoll/scattering_routines/geometry/rotation.h +23 -0
  42. xcoll/scattering_routines/geometry/segments.h +226 -0
  43. xcoll/scattering_routines/geometry/sort.h +184 -0
  44. {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/METADATA +1 -1
  45. {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/RECORD +48 -33
  46. xcoll/beam_elements/collimators_src/absorber.h +0 -141
  47. xcoll/collimator_settings.py +0 -457
  48. xcoll/impacts/__init__.py +0 -1
  49. xcoll/impacts/impacts.py +0 -102
  50. xcoll/impacts/impacts_src/impacts.h +0 -99
  51. xcoll/scattering_routines/everest/crystal.h +0 -1302
  52. xcoll/scattering_routines/everest/scatter.h +0 -169
  53. xcoll/scattering_routines/everest/scatter_crystal.h +0 -260
  54. {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/LICENSE +0 -0
  55. {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/NOTICE +0 -0
  56. {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/WHEEL +0 -0
@@ -1,457 +0,0 @@
1
- import io
2
- import json
3
- import numpy as np
4
- import pandas as pd
5
-
6
-
7
- # in colldb:
8
- # (gap [sigma] + offset [m] + tilt [deg]) or opening [m]
9
- # angle
10
- # length (all three)
11
- # side
12
- # parking ??
13
- # material
14
- # stage
15
- # family
16
- # crystal
17
-
18
- # in element:
19
- # jaw
20
- # ref
21
- # angle
22
- # length (all three)
23
- # side
24
- # material
25
- # is_active
26
-
27
- # in neither (as they are sequence/session-dependent):
28
- # align_to
29
- # s_center
30
- # collimator_type
31
-
32
- # if physical_opening is used, other variables like tilt, opening, offset, .. are ignored
33
- _element_properties = {
34
- 'active_length': 0,
35
- 'inactive_front': 0,
36
- 'inactive_back': 0,
37
- 'jaw': None,
38
- 'reference_center': None,
39
- 'angle': 0,
40
- 'side': 'both',
41
- 'material': None,
42
- 'active': True,
43
- 'crystal': False
44
- }
45
- _crystal_properties = {
46
- 'bend': None,
47
- 'xdim': 0,
48
- 'ydim': 0,
49
- 'miscut': 0,
50
- 'thick': 0,
51
- # 'align_angle': xo.Float64, # = - sqrt(eps/beta)*alpha*nsigma
52
- # 'crytilt': xo.Float64,
53
- # 'orient': xo.Float64,
54
- }
55
- _sequence_properties = {
56
- 's_center': None,
57
- 'align_to': None,
58
- 'collimator_type': None,
59
- }
60
- _colldb_properties = {
61
- 'gap': None, # [L, R]
62
- 'offset': 0, # single value: shift of gap in mm
63
- 'extra_mm': 0, # single value: widening of gap
64
- 'opening_mm': None, # [L, R]
65
- 'tilt': 0,
66
- 'parking': 1,
67
- 'stage': None,
68
- 'family': None,
69
- 'overwritten_keys': []
70
- }
71
- _optics_vals = ['x', 'px', 'y', 'py', 'betx', 'bety', 'alfx', 'alfy', 'dx', 'dy', 'sigma_x', 'sigma_y']
72
-
73
-
74
- # This creates a view on the settings of one collimator, kept in sync with the main database
75
- class CollimatorSettings:
76
-
77
- def __init__(self, name, colldb=None, optics=None, element=None):
78
- if colldb is None:
79
- colldb = pd.DataFrame({**_element_properties, **_crystal_properties}, index=[name])
80
- self._colldb = colldb
81
- if name not in colldb.index:
82
- raise ValueError(f"Collimator {name} not found in database!")
83
- self._optics = optics
84
- self._element = element
85
- self._name = name
86
- self._jaws_manually_set = False
87
-
88
- # Automatically assign all properties from _collimator_properties
89
- # Now create @property's for each of them:
90
- for pp in set(_coll_properties.keys()).union(set(_crystal_properties.keys())):
91
- if not pp in self._colldb.columns:
92
- raise ValueError(f"Error in settings for {self.name}: "
93
- + f"Could not find property {pp} in database column header!")
94
- if pp in _properties_no_setter:
95
- setattr(self.__class__, pp, property(_prop_fget(self, pp)))
96
- else:
97
- setattr(self.__class__, pp, property(_prop_fget(self, pp), _prop_fset(self, pp)))
98
-
99
- @property
100
- def name(self):
101
- return self._name
102
-
103
- def to_dict(self):
104
- props = list(_coll_properties.keys()) + list(_crystal_properties.keys()) + _add_to_dict
105
- all_properties = {pp: getattr(self,pp) for pp in props if pp not in _properties_no_setter}
106
- return {'name': self.name, **all_properties}
107
-
108
- # Some easy accessors to the LR / LRUD properties:
109
- # -----------------------------------------------
110
-
111
- @property
112
- def angle(self):
113
- return _get_LR(self, 'angle')
114
-
115
- @angle.setter
116
- def angle(self, val):
117
- _set_LR(self, 'angle', val)
118
- self._compute_jaws()
119
-
120
- @property
121
- def tilt(self):
122
- return _get_LR(self, 'tilt')
123
-
124
- @tilt.setter
125
- def tilt(self, val):
126
- _set_LR(self, 'tilt', val)
127
- self._compute_jaws()
128
-
129
- @property
130
- def opening(self):
131
- if self._jaws_manually_set:
132
- # TODO: update gap using beam size as from jaws
133
- pass
134
- else:
135
- return _get_LR(self, 'gap')
136
-
137
- @opening.setter
138
- def opening(self, val):
139
- _set_LR(self, 'gap', val, neg=True)
140
- self._jaws_manually_set = False
141
- self._compute_jaws()
142
-
143
- @property
144
- def physical_opening(self):
145
- return _get_LRUD(self, 'jaw', neg=True)
146
-
147
- @physical_opening.setter
148
- def physical_opening(self, val):
149
- _set_LRUD(self, 'jaw', val, neg=True)
150
- self._jaws_manually_set = True
151
- self._compute_jaws()
152
-
153
- @property
154
- def reference_center(self):
155
- return _get_LRUD(self, 'ref', name_LU='_xU', name_RU='_yU', name_LD='_xD', name_RD='_yD')
156
-
157
- @property
158
- def beam_size(self):
159
- if not self._optics_is_ready:
160
- return None
161
- else:
162
- # TODO: calculate beam size from optic
163
- # [[LU,RU], [LC, RC], [LD,RD]]
164
- pass
165
-
166
-
167
- @property
168
- def _optics_is_ready(self):
169
- if self.align_to is None or self._optics is None:
170
- return False
171
- # TODO: when True
172
-
173
- def _compute_jaws(self):
174
- if self._optics_is_ready:
175
- if self._jaws_manually_set:
176
- jaw_LU = self.jaw_LU
177
- jaw_RU = self.jaw_RU
178
- jaw_LD = self.jaw_LD
179
- jaw_RD = self.jaw_RD
180
- # TODO: upate gap
181
- pass
182
- else:
183
- # Default to parking
184
- # TODO: parking is wrt CO, need to correct
185
- jaw_LU = self.parking
186
- jaw_RU = -self.parking
187
- jaw_LD = self.parking
188
- jaw_RD = -self.parking
189
-
190
- if self.gap_L is not None or self.gap_R is not None:
191
- # Get the beam size to be used, depending on align_to
192
- if self.align_to == 'maximum':
193
- # Reset align_to to the location of the maximum
194
- self.align_to = ['front', 'center', 'back'][np.argmax(self.beam_size)]
195
- if self.align_to == 'angular':
196
- sigma_U = self.beam_size[0]
197
- sigma_D = self.beam_size[2]
198
- elif self.align_to == 'front':
199
- sigma_U = self.beam_size[0]
200
- sigma_D = self.beam_size[0]
201
- elif self.align_to == 'center':
202
- sigma_U = self.beam_size[1]
203
- sigma_D = self.beam_size[1]
204
- elif self.align_to == 'back':
205
- sigma_U = self.beam_size[2]
206
- sigma_D = self.beam_size[2]
207
-
208
- # Get the shift due to the tilt to be used, depending on align_to
209
- # TODO: is this correct?
210
- # TODO: better name for align_to?
211
- if align_to == front:
212
- scale = 0
213
- elif align_to == back:
214
- scale = 1
215
- else:
216
- scale = 0.5
217
- ts_LU = -np.tan(self.tilt_L)*self.active_length*scale
218
- ts_RU = -np.tan(self.tilt_R)*self.active_length*scale
219
- ts_LD = np.tan(self.tilt_L)*self.active_length*(1-scale)
220
- ts_RD = np.tan(self.tilt_R)*self.active_length*(1-scale)
221
-
222
- if self.side in ['left', 'both']:
223
- jaw_LU = self.gap_L*sigma_U[0] + self.offset + ts_LU
224
- jaw_LD = self.gap_L*sigma_D[0] + self.offset + ts_LD
225
- if self.side in ['right', 'both']:
226
- jaw_RU = -self.gap_R*sigma_U[1] + self.offset + ts_RU
227
- jaw_RD = -self.gap_R*sigma_D[1] + self.offset + ts_RD
228
-
229
- if self.side == 'left':
230
- self.jaw_LU = min(jaw_LU, self.parking)
231
- self.jaw_LD = min(jaw_LD, self.parking)
232
- self.jaw_RU = None
233
- self.jaw_RD = None
234
- elif self.side == 'right':
235
- self.jaw_LU = None
236
- self.jaw_LD = None
237
- self.jaw_RU = max(jaw_RU, -self.parking)
238
- self.jaw_RD = max(jaw_RD, -self.parking)
239
- else:
240
- self.jaw_LU = min(jaw_LU, self.parking)
241
- self.jaw_LD = min(jaw_LD, self.parking)
242
- self.jaw_RU = max(jaw_RU, -self.parking)
243
- self.jaw_RD = max(jaw_RD, -self.parking)
244
-
245
-
246
- # Helper functions to set/get properties
247
- # --------------------------------------
248
-
249
- def _get_LR(obj, prop, neg=False, name_L='_L', name_R='_R'):
250
- # Is the property reflected along left/right?
251
- sign = -1 if neg else 1
252
- # Is it a property or a dict key?
253
- if isinstance(obj, dict):
254
- L = obj[prop + name_L]
255
- R = obj[prop + name_R]
256
- else:
257
- L = getattr(obj, prop + name_L)
258
- R = getattr(obj, prop + name_R)
259
- # Find out how many values to return
260
- if L is None and R is None:
261
- return None
262
- elif L is None:
263
- return R
264
- elif R is None:
265
- return L
266
- else:
267
- return L if L==sign*R else [L,R]
268
-
269
-
270
- def _set_LR(obj, prop, val, neg=False, name=None, name_L='_L', name_R='_R'):
271
- # 'name' is only used for error reporting
272
- if name is None:
273
- if isinstance(obj, dict):
274
- name = 'dict_property'
275
- else:
276
- name = obj.name if hasattr(obj, 'name') else obj.__class__.__name__
277
- # Is the property reflected along left/right?
278
- sign = -1 if neg else 1
279
- # Find out how to set values
280
- if not hasattr(val, '__iter__'):
281
- val = [val]
282
- error_dimension = False
283
- if isinstance(val, str):
284
- raise ValueError(f"Error in settings for {name}: "
285
- + f"The setting `{prop}` has to be a number!")
286
- elif len(val) == 2:
287
- # The value is of the form [val_L,val_R]
288
- if hasattr(val[0], '__iter__') or hasattr(val[1], '__iter__'):
289
- error_dimension = True
290
- val_L = val[0]
291
- val_R = val[1]
292
- elif len(val) == 1:
293
- # The value is of the form [val]
294
- if hasattr(val[0], '__iter__'):
295
- error_dimension = True
296
- val_L = val[0]
297
- val_R = sign*val[0]
298
- else:
299
- error_dimension = True
300
- if error_dimension:
301
- raise ValueError(f"Error in settings for {name}: "
302
- + f"The setting `{prop}` must have one or two (L, R) values!")
303
- # Is it a property or a dict key?
304
- if isinstance(obj, dict):
305
- obj[prop + name_L] = val_L
306
- obj[prop + name_R] = val_R
307
- else:
308
- _prop_fset(obj, prop + name_L)(obj, val_L)
309
- _prop_fset(obj, prop + name_R)(obj, val_R)
310
-
311
-
312
- def _get_LRUD(obj, prop, neg=False, name=None,
313
- name_LU='_LU', name_RU='_RU', name_LD='_LD', name_RD='_RD'):
314
- # 'name' is only used for error reporting
315
- if name is None:
316
- if isinstance(obj, dict):
317
- name = 'dict_property'
318
- else:
319
- name = obj.name if hasattr(obj, 'name') else obj.__class__.__name__
320
- # Is the property reflected along left/right?
321
- sign = -1 if neg else 1
322
- # Is it a property or a dict key?
323
- if isinstance(obj, dict):
324
- LU = obj[prop + name_LU]
325
- RU = obj[prop + name_RU]
326
- LD = obj[prop + name_LD]
327
- RD = obj[prop + name_RD]
328
- else:
329
- LU = getattr(obj, prop + name_LU)
330
- RU = getattr(obj, prop + name_RU)
331
- LD = getattr(obj, prop + name_LD)
332
- RD = getattr(obj, prop + name_RD)
333
- # Find out how many values to return
334
- if LU is None and RU is None \
335
- and LD is None and RD is None:
336
- return None
337
- elif LU is None and RU is None:
338
- return LD if LD==sign*RD else [LD,RD]
339
- elif LD is None and RD is None:
340
- return LU if LU==sign*RU else [LU,RU]
341
- elif LU is None and RD is None:
342
- raise ValueError(f"Error in settings for {name}: "
343
- + f"The setting `{prop}` has values for LD and RU but not for LU and RD."
344
- + f"Cannot mix jaws L/R with U/D! Either set all four, or only L or only R.")
345
- elif LD is None and RU is None:
346
- raise ValueError(f"Error in settings for {name}: "
347
- + f"The setting `{prop}` has values for LU and RD but not for LD and RU."
348
- + f"Cannot mix jaws L/R with U/D! Either set all four, or only L or only R.")
349
- else:
350
- if LU == sign*RU == LD == sign*RD:
351
- return LU
352
- elif LU == LD and RU == RD:
353
- return [LU, RU]
354
- else:
355
- return [[LU, RU], [LD, RD]]
356
-
357
-
358
- def _set_LRUD(obj, prop, val, neg=False, name=None,
359
- name_LU='_LU', name_RU='_RU', name_LD='_LD', name_RD='_RD'):
360
- # 'name' is only used for error reporting
361
- if name is None:
362
- if isinstance(obj, dict):
363
- name = 'dict_property'
364
- else:
365
- name = obj.name if hasattr(obj, 'name') else obj.__class__.__name__
366
- # Is the property reflected along left/right?
367
- sign = -1 if neg else 1
368
-
369
- # Find out how to set values
370
- if not hasattr(val, '__iter__'):
371
- val = [val]
372
- error_string = False
373
- error_dimension = False
374
- if isinstance(val, str):
375
- error_string = True
376
- elif len(val) == 4:
377
- # The value is of the form [val_LU,val_RU,val_LD,val_RD]
378
- if hasattr(val[0], '__iter__') or hasattr(val[1], '__iter__') \
379
- or hasattr(val[2], '__iter__') or hasattr(val[3], '__iter__'):
380
- error_dimension = True
381
- val_LU = val[0]
382
- val_RU = val[1]
383
- val_LD = val[2]
384
- val_RD = val[3]
385
- elif len(val) == 2:
386
- if not hasattr(val[0], '__iter__') and not hasattr(val[1], '__iter__'):
387
- # The value is of the form [val_L,val_R]
388
- val_LU = val[0]
389
- val_RU = val[1]
390
- val_LD = val[0]
391
- val_RD = val[1]
392
- elif not hasattr(val[0], '__iter__') or not hasattr(val[1], '__iter__'):
393
- error_dimension = True
394
- else:
395
- # The value is of the form [[val_LU, val_RU], [val_LD, val_RD]]
396
- if isinstance(val[0], str) or isinstance(val[1], str):
397
- error_string = True
398
- if hasattr(val[0][0], '__iter__') or hasattr(val[0][1], '__iter__') \
399
- or hasattr(val[1][0], '__iter__') or hasattr(val[1][1], '__iter__'):
400
- error_dimension = True
401
- val_LU = val[0][0]
402
- val_RU = val[0][1]
403
- val_LD = val[1][0]
404
- val_RD = val[1][1]
405
- elif len(val) == 1:
406
- # The value is of the form [val]
407
- if hasattr(val[0], '__iter__'):
408
- error_dimension = True
409
- val_LU = val[0]
410
- val_RU = sign*val[0]
411
- val_LD = val[0]
412
- val_RD = sign*val[0]
413
- else:
414
- error_dimension = True
415
- if error_dimension:
416
- raise ValueError(f"Error in settings for {name}: "
417
- + f"The setting `{prop}` has to be given as val, [val], "
418
- + f"[val_L, val_R], [[val_LU, val_RU], [val_LD, val_RD]], "
419
- + f"or [val_LU, val_RU, val_LD, val_RD]!")
420
- if error_string:
421
- raise ValueError(f"Error in settings for {name}: "
422
- + f"The setting `{prop}` has to be a number or a "
423
- + f"list of numbers!")
424
-
425
- # Is it a property or a dict key?
426
- if isinstance(obj, dict):
427
- obj[prop + name_LU] = val_LU
428
- obj[prop + name_RU] = val_RU
429
- obj[prop + name_LD] = val_LD
430
- obj[prop + name_RD] = val_RD
431
- else:
432
- _prop_fset(obj, prop + name_LU)(obj, val_LU)
433
- _prop_fset(obj, prop + name_RU)(obj, val_RU)
434
- _prop_fset(obj, prop + name_LD)(obj, val_LD)
435
- _prop_fset(obj, prop + name_RD)(obj, val_RD)
436
-
437
-
438
- # These getter and setter functions link each property to the corresponding
439
- # entry in the DataFrame
440
- def _prop_fget(obj, attr_name):
441
- def fget(obj):
442
- return obj._colldb.loc[obj.name, attr_name]
443
- return fget
444
-
445
- def _prop_fset(obj, attr_name):
446
- def fset(obj, value):
447
- if hasattr(obj, '_colldb'):
448
- obj._colldb.loc[obj.name, attr_name] = value
449
- if hasattr(obj, attr_name):
450
- setattr(obj, attr_name, value)
451
- # If we want additional effects on the setter funcions, this
452
- # can be achieved by defining an fset_prop method
453
- if hasattr(obj.__class__, 'fset_' + attr_name):
454
- getattr(obj, 'fset_' + attr_name)(value)
455
- # TODO: self._compute_jaws()
456
- return fset
457
-
xcoll/impacts/__init__.py DELETED
@@ -1 +0,0 @@
1
- from .impacts import CollimatorImpacts
xcoll/impacts/impacts.py DELETED
@@ -1,102 +0,0 @@
1
- # copyright ############################### #
2
- # This file is part of the Xcoll Package. #
3
- # Copyright (c) CERN, 2023. #
4
- # ######################################### #
5
-
6
- import xobjects as xo
7
- import xtrack as xt
8
-
9
- from .interaction_types import source, interactions, shortcuts
10
- from ..general import _pkg_root
11
-
12
- import numpy as np
13
- import pandas as pd
14
-
15
-
16
- class CollimatorImpacts(xt.BeamElement):
17
- _xofields = {
18
- '_index': xt.RecordIndex,
19
- 'at_element': xo.Int64[:],
20
- 'at_turn': xo.Int64[:],
21
- 'ds': xo.Float64[:],
22
- '_inter': xo.Int64[:],
23
- 'parent_id': xo.Int64[:],
24
- 'parent_x': xo.Float64[:],
25
- 'parent_px': xo.Float64[:],
26
- 'parent_y': xo.Float64[:],
27
- 'parent_py': xo.Float64[:],
28
- 'parent_zeta': xo.Float64[:],
29
- 'parent_delta': xo.Float64[:],
30
- 'parent_energy': xo.Float64[:],
31
- 'parent_mass': xo.Float64[:],
32
- 'parent_charge': xo.Int64[:],
33
- 'parent_z': xo.Int64[:],
34
- 'parent_a': xo.Int64[:],
35
- 'parent_pdgid': xo.Int64[:],
36
- 'child_id': xo.Int64[:],
37
- 'child_x': xo.Float64[:],
38
- 'child_px': xo.Float64[:],
39
- 'child_y': xo.Float64[:],
40
- 'child_py': xo.Float64[:],
41
- 'child_zeta': xo.Float64[:],
42
- 'child_delta': xo.Float64[:],
43
- 'child_energy': xo.Float64[:],
44
- 'child_mass': xo.Float64[:],
45
- 'child_charge': xo.Int64[:],
46
- 'child_z': xo.Int64[:],
47
- 'child_a': xo.Int64[:],
48
- 'child_pdgid': xo.Int64[:],
49
- }
50
-
51
- allow_track = False
52
-
53
- _extra_c_sources = [
54
- source,
55
- _pkg_root.joinpath('headers','particle_states.h'),
56
- _pkg_root.joinpath('impacts','impacts_src','impacts.h')
57
- ]
58
-
59
- @property
60
- def interaction_type(self):
61
- return np.array([interactions[inter] for inter in self._inter])
62
-
63
- def collimator_name(self, element_id):
64
- if not hasattr(self, '_coll_ids'):
65
- return element_id
66
- elif element_id not in self._coll_ids:
67
- raise ValueError(f"Element {element_id} not found in list of collimators of this impact table!\n"
68
- + "Did the line change without updating the list in the impact table?")
69
- else:
70
- return self._coll_ids[element_id]
71
-
72
- def to_pandas(self):
73
- n_rows = self._index.num_recorded
74
- coll_header = 'collimator' if hasattr(self, '_coll_ids') else 'collimator id'
75
- df = pd.DataFrame({
76
- 'turn': self.at_turn[:n_rows],
77
- coll_header: [self.collimator_name(element_id) for element_id in self.at_element[:n_rows]],
78
- 'interaction_type': [interactions[inter] for inter in self._inter[:n_rows]],
79
- 'ds': self.ds[:n_rows],
80
- **{
81
- f'{p}_{val}': getattr(self, f'{p}_{val}')[:n_rows]
82
- for p in ['parent', 'child']
83
- for val in ['id', 'x', 'px', 'y', 'py', 'zeta', 'delta', 'energy', 'mass', 'charge', 'z', 'a', 'pdgid']
84
- }
85
- })
86
- return df
87
-
88
- # TODO: list of impacted collimators
89
-
90
-
91
- # TODO: does not work when multiple children
92
- # TODO: allow to select collimator by name
93
- def interactions_per_collimator(self, collimator_id=0, *, turn=None):
94
- mask = (self._inter > 0) & (self.at_element == collimator_id)
95
- if turn is not None:
96
- mask = mask & (self.at_turn == turn)
97
- df = pd.DataFrame({
98
- 'int': [shortcuts[inter] for inter in self._inter[mask]],
99
- 'pid': self.parent_id[mask]
100
- })
101
- return df.groupby('pid', sort=False)['int'].agg(list)
102
-
@@ -1,99 +0,0 @@
1
- // copyright ############################### #
2
- // This file is part of the Xcoll Package. #
3
- // Copyright (c) CERN, 2023. #
4
- // ######################################### #
5
-
6
-
7
- #ifndef XCOLL_IMPACTS_H
8
- #define XCOLL_IMPACTS_H
9
-
10
- // TODO: do we need to pass RecordIndex?
11
- // probably can do RecordIndex record_index = CollimatorImpactsData_getp__index(record); ?
12
- /*gpufun*/
13
- int64_t CollimatorImpactsData_log(CollimatorImpactsData record, RecordIndex record_index, LocalParticle* parent,
14
- int64_t interaction){
15
- // This can be used for a point-like interaction where there is no child (or because it's equal to the parent)
16
- // or to log the parent first, to be followed up with CollimatorImpactsData_log_child on the same slot
17
-
18
- int64_t i_slot = -1;
19
- if (record){
20
- // Get a slot in the record (this is thread safe)
21
- i_slot = RecordIndex_get_slot(record_index);
22
- // The returned slot id is negative if record is NULL or if record is full
23
-
24
- if (i_slot>=0){
25
- CollimatorImpactsData_set_at_element(record, i_slot, LocalParticle_get_at_element(parent));
26
- CollimatorImpactsData_set_at_turn(record, i_slot, LocalParticle_get_at_turn(parent));
27
- CollimatorImpactsData_set_ds(record, i_slot, 0);
28
- CollimatorImpactsData_set__inter(record, i_slot, interaction);
29
-
30
- double charge_ratio = LocalParticle_get_charge_ratio(parent);
31
- double mass_ratio = charge_ratio / LocalParticle_get_chi(parent);
32
- double energy = ( LocalParticle_get_ptau(parent) + 1 / LocalParticle_get_beta0(parent)
33
- ) * mass_ratio * LocalParticle_get_p0c(parent);
34
- // All fields have to be written, or the arrays will not have the same length
35
- // TODO: maybe this is not true, as we are setting by slot index? Don't the arrays come pre-initialised?
36
- CollimatorImpactsData_set_parent_id(record, i_slot, LocalParticle_get_particle_id(parent));
37
- CollimatorImpactsData_set_parent_x(record, i_slot, LocalParticle_get_x(parent));
38
- CollimatorImpactsData_set_parent_px(record, i_slot, LocalParticle_get_px(parent));
39
- CollimatorImpactsData_set_parent_y(record, i_slot, LocalParticle_get_y(parent));
40
- CollimatorImpactsData_set_parent_py(record, i_slot, LocalParticle_get_py(parent));
41
- CollimatorImpactsData_set_parent_zeta(record, i_slot, LocalParticle_get_zeta(parent));
42
- CollimatorImpactsData_set_parent_delta(record, i_slot, LocalParticle_get_delta(parent));
43
- CollimatorImpactsData_set_parent_energy(record, i_slot, energy);
44
- CollimatorImpactsData_set_parent_mass(record, i_slot, mass_ratio*LocalParticle_get_mass0(parent));
45
- CollimatorImpactsData_set_parent_charge(record, i_slot, charge_ratio*LocalParticle_get_q0(parent));
46
- // TODO: particle info
47
- CollimatorImpactsData_set_parent_z(record, i_slot, -1);
48
- CollimatorImpactsData_set_parent_a(record, i_slot, -1);
49
- CollimatorImpactsData_set_parent_pdgid(record, i_slot, -1);
50
-
51
- // TODO: maybe this is not needed
52
- CollimatorImpactsData_set_child_id(record, i_slot, -1);
53
- CollimatorImpactsData_set_child_x(record, i_slot, -1);
54
- CollimatorImpactsData_set_child_px(record, i_slot, -1);
55
- CollimatorImpactsData_set_child_y(record, i_slot, -1);
56
- CollimatorImpactsData_set_child_py(record, i_slot, -1);
57
- CollimatorImpactsData_set_child_zeta(record, i_slot, -1);
58
- CollimatorImpactsData_set_child_delta(record, i_slot, -1);
59
- CollimatorImpactsData_set_child_energy(record, i_slot, -1);
60
- CollimatorImpactsData_set_child_mass(record, i_slot, -1);
61
- CollimatorImpactsData_set_child_charge(record, i_slot, -1);
62
- CollimatorImpactsData_set_child_z(record, i_slot, -1);
63
- CollimatorImpactsData_set_child_a(record, i_slot, -1);
64
- CollimatorImpactsData_set_child_pdgid(record, i_slot, -1);
65
- }
66
- }
67
- // printf("Logging %i in slot %i\n", interaction, i_slot);
68
- return i_slot;
69
- }
70
-
71
- /*gpufun*/
72
- void CollimatorImpactsData_log_child(CollimatorImpactsData record, int64_t i_slot, LocalParticle* child, double ds){
73
-
74
- if (record && i_slot>=0){
75
- CollimatorImpactsData_set_ds(record, i_slot, ds);
76
-
77
- double charge_ratio = LocalParticle_get_charge_ratio(child);
78
- double mass_ratio = charge_ratio / LocalParticle_get_chi(child);
79
- double energy = ( LocalParticle_get_ptau(child) + 1 / LocalParticle_get_beta0(child)
80
- ) * mass_ratio * LocalParticle_get_p0c(child);
81
- CollimatorImpactsData_set_child_id(record, i_slot, LocalParticle_get_particle_id(child));
82
- CollimatorImpactsData_set_child_x(record, i_slot, LocalParticle_get_x(child));
83
- CollimatorImpactsData_set_child_px(record, i_slot, LocalParticle_get_px(child));
84
- CollimatorImpactsData_set_child_y(record, i_slot, LocalParticle_get_y(child));
85
- CollimatorImpactsData_set_child_py(record, i_slot, LocalParticle_get_py(child));
86
- CollimatorImpactsData_set_child_zeta(record, i_slot, LocalParticle_get_zeta(child));
87
- CollimatorImpactsData_set_child_delta(record, i_slot, LocalParticle_get_delta(child));
88
- CollimatorImpactsData_set_child_energy(record, i_slot, energy);
89
- CollimatorImpactsData_set_child_mass(record, i_slot, mass_ratio*LocalParticle_get_mass0(child));
90
- CollimatorImpactsData_set_child_charge(record, i_slot, charge_ratio*LocalParticle_get_q0(child));
91
- // TODO: particle info
92
- CollimatorImpactsData_set_child_z(record, i_slot, -1);
93
- CollimatorImpactsData_set_child_a(record, i_slot, -1);
94
- CollimatorImpactsData_set_child_pdgid(record, i_slot, -1);
95
- // printf("Slot %i: length %f\n", i_slot, ds);
96
- }
97
- }
98
-
99
- #endif /* XCOLL_IMPACTS_H */