xcoll 0.3.5__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.5.dist-info → xcoll-0.4.0.dist-info}/METADATA +1 -1
  45. {xcoll-0.3.5.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.5.dist-info → xcoll-0.4.0.dist-info}/LICENSE +0 -0
  55. {xcoll-0.3.5.dist-info → xcoll-0.4.0.dist-info}/NOTICE +0 -0
  56. {xcoll-0.3.5.dist-info → xcoll-0.4.0.dist-info}/WHEEL +0 -0
xcoll/colldb.py CHANGED
@@ -1,6 +1,6 @@
1
1
  # copyright ############################### #
2
2
  # This file is part of the Xcoll Package. #
3
- # Copyright (c) CERN, 2023. #
3
+ # Copyright (c) CERN, 2024. #
4
4
  # ######################################### #
5
5
 
6
6
  import io
@@ -8,26 +8,22 @@ import json
8
8
  import numpy as np
9
9
  import pandas as pd
10
10
 
11
+ import xtrack as xt
11
12
 
12
- #TODO: niet-active collimators op non-active etc (ook crystals)
13
+ from .beam_elements import BlackAbsorber, EverestCollimator, EverestCrystal, collimator_classes, element_classes
14
+ from .install import install_elements
15
+ from .scattering_routines.everest.materials import SixTrack_to_xcoll
13
16
 
14
- def load_SixTrack_colldb(filename, *, emit):
15
- print("Warning: Using 'xcoll.load_SixTrack_colldb()' is deprecated! "
16
- + "Use 'xcoll.CollimatorDatabase.from_SixTrack()' instead.")
17
- return CollimatorDatabase.from_SixTrack(file=filename, nemitt_x=emit, nemitt_y=emit)
18
17
 
19
-
20
- def _initialise_None(collimator):
21
- fields = {'s_center':None, 'align_to': None, 's_align_front': None, 's_align_back': None }
22
- fields.update({'gap_L': None, 'gap_R': None, 'angle_L': 0, 'angle_R': 0, 'offset': 0, 'parking': 1})
23
- fields.update({'jaw_LU': None, 'jaw_RU': None, 'jaw_LD': None, 'jaw_RD': None, 'family': None, 'overwritten_keys': []})
24
- fields.update({'side': 'both', 'material': None, 'stage': None, 'collimator_type': None, 'active': True})
25
- fields.update({'active_length': 0, 'inactive_front': 0, 'inactive_back': 0, 'sigmax': None, 'sigmay': None})
26
- fields.update({'crystal': None, 'bending_radius': None, 'xdim': 0, 'ydim': 0, 'miscut': 0, 'thick': 0})
18
+ def _initialise_None(dct):
19
+ fields = {'gap': None, 'angle': 0, 'offset': 0, 'parking': 1, 'jaw': None, 'family': None}
20
+ fields.update({'overwritten_keys': [], 'side': 'both', 'material': None, 'stage': None})
21
+ fields.update({'length': 0, 'collimator_type': None, 'active': True, 'crystal': None, 'tilt': 0})
22
+ fields.update({'bending_radius': 0, 'bending_angle': 0, 'width': 0, 'height': 0, 'miscut': 0})
27
23
  for f, val in fields.items():
28
- if f not in collimator.keys():
29
- collimator[f] = val
30
- for key in collimator.keys():
24
+ if f not in dct.keys():
25
+ dct[f] = val
26
+ for key in dct.keys():
31
27
  if key not in fields.keys():
32
28
  raise ValueError(f"Illegal setting {key} in collimator!")
33
29
 
@@ -47,7 +43,7 @@ def _get_coll_dct_by_beam(coll, beam):
47
43
  beam = f'b{beam}'
48
44
  beam = beam.lower()
49
45
  beam_in_db = list(coll.keys())
50
-
46
+
51
47
  if beam_in_db == ['b1','b2']:
52
48
  if beam is None:
53
49
  raise ValueError("Need to specify a beam, because the given dict is for both beams!")
@@ -67,13 +63,88 @@ def _get_coll_dct_by_beam(coll, beam):
67
63
 
68
64
  class CollimatorDatabase:
69
65
 
70
- _init_vars = ['collimator_dict', 'family_dict', 'beam', 'nemitt_x', 'nemitt_y', '_yaml_merged', 'ignore_crystals']
71
- _init_var_defaults = {'family_dict': {}, 'beam': None, '_yaml_merged': False, 'ignore_crystals': True}
66
+ def __init__(self, **kwargs):
67
+ # Get all arguments
68
+ for key in ['collimator_dict', 'nemitt_x', 'nemitt_y']:
69
+ if key not in kwargs.keys():
70
+ raise ValueError(f"CollimatorDatabase is missing required argument '{key}'!")
71
+
72
+ self._parse_dict(kwargs['collimator_dict'],
73
+ kwargs.get('family_dict', {}),
74
+ kwargs.get('beam', None),
75
+ kwargs.get('_yaml_merged', False),
76
+ kwargs.get('ignore_crystals', True))
77
+ self.nemitt_x = kwargs['nemitt_x']
78
+ self.nemitt_y = kwargs['nemitt_y']
79
+ self.line = None
80
+
81
+
82
+ def _parse_dict(self, coll, fam, beam=None, _yaml_merged=False, ignore_crystals=True):
83
+ # We make all keys case-insensitive to avoid confusion between different conventions
84
+ coll = _dict_keys_to_lower(coll)
85
+ fam = _dict_keys_to_lower(fam)
86
+
87
+ # The dictionary can be a CollimatorDatabase for a single beam (beam=None)
88
+ # or for both beams (beam='b1' or beam='b2)
89
+ coll = _get_coll_dct_by_beam(coll, beam)
90
+
91
+ # Apply family settings
92
+ crystals = []
93
+ for thiscoll, settings in coll.items():
94
+ settings = {k.lower(): v for k,v in settings.items()}
95
+ if 'family' in settings.keys() and settings['family'] is not None:
96
+ settings['family'] = settings['family'].lower()
97
+ thisfam = settings['family']
98
+ if thisfam not in fam.keys():
99
+ raise ValueError(f"Collimator {thiscoll} depends on family {thisfam}, "
100
+ + f"but the latter is not defined!")
101
+
102
+ # Check if some family settings are overwritten for this collimator
103
+ # Only do this check if we didn't do a YAML merge earlier (because then it
104
+ # is already taken care of)
105
+ if not _yaml_merged:
106
+ overwritten_keys = [key.lower() for key in settings.keys() if key in fam[thisfam]]
107
+ if len(overwritten_keys) > 0:
108
+ settings['overwritten_keys'] = overwritten_keys
109
+
110
+ # Load family settings, potentially overwriting settings for this collimator
111
+ settings = {**fam[thisfam], **settings}
112
+
113
+ else:
114
+ settings['family'] = None
115
+ coll[thiscoll] = settings
116
+
117
+ # Save list of crystals
118
+ if 'crystal' in settings:
119
+ if settings['crystal'] != 0.0:
120
+ crystals += [thiscoll]
121
+ else:
122
+ settings['crystal'] = None
123
+
124
+ # Remove crystals from colldb
125
+ if ignore_crystals:
126
+ for thiscoll in crystals:
127
+ del coll[thiscoll]
128
+
129
+ # Check that all collimators have gap settings
130
+ if not np.all(['gap' in val.keys() or 'jaw' in val.keys() for val in coll.values()]):
131
+ raise ValueError("Ill-defined CollimatorDatabase: Not all collimators have a gap or "
132
+ + "jaw setting, (or the keys / structure of the dictionary is wrong)!")
133
+
134
+ # Update collimators with default values for missing keys
135
+ for name, collimator in coll.items():
136
+ # Change all values to lower case
137
+ for key, val in collimator.items():
138
+ collimator[key] = val.lower() if isinstance(val, str) else val
139
+ _initialise_None(collimator)
140
+
141
+ self._collimator_dict = coll
142
+ self._family_dict = fam
143
+
144
+ # =======================================
145
+ # ====== Loading/dumping functions ======
146
+ # =======================================
72
147
 
73
- # -------------------------------
74
- # ------ Loading functions ------
75
- # -------------------------------
76
-
77
148
  @classmethod
78
149
  def from_yaml(cls, file, **kwargs):
79
150
 
@@ -119,6 +190,8 @@ class CollimatorDatabase:
119
190
  coll_dct = _get_coll_dct_by_beam(dct['collimators'], beam)
120
191
  full_coll_dct = _get_coll_dct_by_beam(full_dct[collkey], beam)
121
192
  for coll, full_coll in zip(coll_dct.values(), full_coll_dct.values()):
193
+ if not isinstance(coll['gap'], (int,float)):
194
+ coll['gap'] = None
122
195
  if 'family' in coll.keys():
123
196
  raise ValueError(f"Error in {coll}: Cannot use merging for families "
124
197
  + "and manually specify family as well!")
@@ -193,10 +266,10 @@ class CollimatorDatabase:
193
266
  family_settings = {}
194
267
  family_types = {}
195
268
  side = {}
196
- cry_fields = ['bending_radius', 'xdim', 'ydim', 'thick', 'miscut', 'crystal']
269
+ cry_fields = ['bending_radius', 'width', 'height', 'thick', 'miscut', 'crystal']
197
270
  cry = {}
198
271
 
199
- for l_no, line in enumerate(infile):
272
+ for line in infile:
200
273
  if line.startswith('#'):
201
274
  continue # Comment
202
275
  sline = line.split()
@@ -224,16 +297,20 @@ class CollimatorDatabase:
224
297
 
225
298
  defaults = {}
226
299
  _initialise_None(defaults)
300
+ defaults['thick'] = 0
227
301
 
228
- famdct = {key: {'gap': family_settings[key], 'stage': family_types[key]} for key in family_settings}
229
- names = ['name', 'gap', 'material', 'active_length', 'angle', 'offset']
302
+ famdct = {key: {'gap': None if family_settings[key] > 900 else family_settings[key],
303
+ 'stage': family_types[key]} for key in family_settings}
304
+ names = ['name', 'gap', 'material', 'length', 'angle', 'offset']
230
305
 
231
306
  df = pd.read_csv(io.StringIO(coll_data_string), sep=r'\s+', index_col=False, names=names)
232
307
  df['family'] = df['gap'].copy()
233
308
  df['family'] = df['family'].apply(lambda s: None if re.match(r'^-?\d+(\.\d+)?$', str(s)) else s)
234
- df.insert(5,'stage', df['gap'].apply(lambda s: family_types.get(s, 'UNKNOWN')))
309
+ df.insert(5,'stage', df['gap'].apply(lambda s: None if s in family_types else 'UNKNOWN'))
310
+
311
+ df['gap'] = df['gap'].apply(lambda s: None if not isinstance(s, str) and s > 900 else s)
312
+ df['gap'] = df['gap'].apply(lambda s: None if isinstance(s, str) else s)
235
313
 
236
- df['gap'] = df['gap'].apply(lambda s: float(family_settings.get(s, s)))
237
314
  # TODO this breaks code if a key has upper case, e.g. gap_L
238
315
  df['name'] = df['name'].str.lower() # Make the names lowercase for easy processing
239
316
  df['parking'] = 0.025
@@ -243,6 +320,9 @@ class CollimatorDatabase:
243
320
  for key in cry_fields:
244
321
  df[key] = [cry[name][key] if name in cry else defaults[key]
245
322
  for name in df['name']]
323
+ if not np.allclose(np.unique(df.thick.values), 0):
324
+ print("Warning: Keyword 'thick' is currently not supported in xcoll! Ignoring.")
325
+ df = df.drop('thick', axis=1)
246
326
  df['crystal'] = ['strip' if s==1 else s for s in df['crystal']]
247
327
  df['crystal'] = ['quasi-mosaic' if s==2 else s for s in df['crystal']]
248
328
  df['side'] = [side[name] if name in side else defaults['side']
@@ -250,13 +330,32 @@ class CollimatorDatabase:
250
330
  df['side'] = ['both' if s==0 else s for s in df['side']]
251
331
  df['side'] = ['left' if s==1 else s for s in df['side']]
252
332
  df['side'] = ['right' if s==2 else s for s in df['side']]
333
+ if not np.allclose(np.unique(df.offset.values), 0):
334
+ print("Warning: Keyword 'offset' is currently not supported in xcoll! Ignoring.")
335
+ df = df.drop('offset', axis=1)
253
336
  df = df.set_index('name')
254
-
255
- return cls.from_dict({'collimators': df.transpose().to_dict(), 'families': famdct}, \
337
+ df = df.replace(np.nan, None)
338
+
339
+ colldict = df.transpose().to_dict()
340
+ # Remove Nonetype families
341
+ colldict = {coll: {kk: vv for kk, vv in coll_settings.items()
342
+ if kk != 'family' or vv is not None}
343
+ for coll, coll_settings in colldict.items()}
344
+ # Remove None gaps and stages if the collimator is assigned a family (they will be set in _parse_dict)
345
+ colldict = {coll: {kk: vv for kk, vv in coll_settings.items()
346
+ if kk != 'gap' or vv is not None or 'family' not in coll_settings}
347
+ for coll, coll_settings in colldict.items()}
348
+ colldict = {coll: {kk: vv for kk, vv in coll_settings.items()
349
+ if kk != 'stage' or vv is not None or 'family' not in coll_settings}
350
+ for coll, coll_settings in colldict.items()}
351
+
352
+ return cls.from_dict({'collimators': colldict, 'families': famdct}, \
256
353
  ignore_crystals=ignore_crystals, **kwargs)
257
354
 
355
+ def to_pandas(self):
356
+ return pd.DataFrame(self._collimator_dict).transpose()
258
357
 
259
- def write_to_yaml(self, out, lhc_style=True):
358
+ def to_yaml(self, out, lhc_style=True):
260
359
  """
261
360
  Writes a colldb in memory to disk in the yaml format.
262
361
 
@@ -286,18 +385,18 @@ class CollimatorDatabase:
286
385
  formatted_key = f'{key}:'.ljust(key_width)
287
386
  #formatted_values = formatted_values.ljust(key_width)
288
387
  return f"{spacing}{formatted_key} {{ {formatted_values} }}\n"
289
-
388
+
290
389
  def _print_values(keys, dct, file, spacing='', mapping=False):
291
390
  # Writes formatted dictionary entries to a file
292
391
  for key in keys:
293
392
  file.write(_format_dict_entry(key, dct[key], spacing=spacing, mapping=mapping))
294
-
393
+
295
394
  def _print_colls(colls, dcts, beam, file):
296
395
  # Filters and formats collimator data, then writes to a file
297
396
  coll_items_to_print = ['<<','gap','angle','material','active','length','side']
298
397
  file.write(f' {beam}:\n')
299
398
  for coll in colls:
300
- coll_dict = dcts._colldb.transpose().to_dict()[coll]
399
+ coll_dict = dcts.to_pandas().transpose().to_dict()[coll]
301
400
  fam = coll_dict['family']
302
401
  fam_keys = []
303
402
  if fam is not None:
@@ -305,19 +404,19 @@ class CollimatorDatabase:
305
404
  coll_dict = {**{'<<': '*'+fam}, **coll_dict}
306
405
  temp_items_to_print = []
307
406
  if coll_dict['crystal'] and str(coll_dict['crystal'])!='nan':
308
- temp_items_to_print = ['bending_radius','xdim','ydim','miscut','crystal', 'thick']
309
- if coll_dict['angle_L'] == coll_dict['angle_R']:
310
- coll_dict.update({'angle': coll_dict['angle_L']})
311
- else:
312
- temp_items_to_print = temp_items_to_print + ['angle_L','angle_R']
313
- if coll_dict['gap_L'] == coll_dict['gap_R']:
314
- coll_dict.update({'gap': coll_dict['gap_L']})
315
- elif coll_dict['gap_L'] is None and coll_dict['gap_R'] is not None:
316
- coll_dict.update({'gap': coll_dict['gap_R']})
317
- elif coll_dict['gap_L'] is not None and coll_dict['gap_R'] is None:
318
- coll_dict.update({'gap': coll_dict['gap_L']})
319
- else:
320
- temp_items_to_print = temp_items_to_print + ['gap_L','gap_R']
407
+ temp_items_to_print = ['bending_radius','width','height','miscut','crystal']
408
+ # if 'angle_L' in coll_dict and coll_dict['angle_L'] == coll_dict['angle_R']:
409
+ # coll_dict.update({'angle': coll_dict['angle_L']})
410
+ # else:
411
+ # temp_items_to_print = temp_items_to_print + ['angle_L','angle_R']
412
+ # if coll_dict['gap_L'] == coll_dict['gap_R']:
413
+ # coll_dict.update({'gap': coll_dict['gap_L']})
414
+ # elif coll_dict['gap_L'] is None and coll_dict['gap_R'] is not None:
415
+ # coll_dict.update({'gap': coll_dict['gap_R']})
416
+ # elif coll_dict['gap_L'] is not None and coll_dict['gap_R'] is None:
417
+ # coll_dict.update({'gap': coll_dict['gap_L']})
418
+ # else:
419
+ # temp_items_to_print = temp_items_to_print + ['gap_L','gap_R']
321
420
  value = {}
322
421
  overwritten_keys = coll_dict['overwritten_keys']
323
422
  for key, val in coll_dict.items():
@@ -326,8 +425,7 @@ class CollimatorDatabase:
326
425
  if (key in coll_items_to_print+temp_items_to_print) and (key not in (set(fam_keys)-set(overwritten_keys))) and (val != 'both'):
327
426
  value.update({key: val})
328
427
  file.write(_format_dict_entry(coll, value, spacing=' '))
329
- file.write('\n')
330
-
428
+ file.write('\n')
331
429
 
332
430
  LHC_families = ['tcp3', 'tcsg3', 'tcsm3', 'tcla3', 'tcp7', 'tcsg7', 'tcsm7', 'tcla7', 'tcli', 'tdi', 'tcdq', 'tcstcdq', 'tcth1', 'tcth2', 'tcth5', 'tcth8', 'tctv1', 'tctv2', 'tctv5', 'tctv8', 'tclp', 'tcxrp', 'tcryo', 'tcl4', 'tcl5', 'tcl6', 'tct15', 'tct2', 'tct8', 'tcsp', 'tcld']
333
431
  with open(f'{out}.yaml', 'w') as file:
@@ -376,15 +474,15 @@ class CollimatorDatabase:
376
474
  _print_values(self._family_dict.keys(), self._family_dict, file, spacing=' - &', mapping=True)
377
475
 
378
476
  # Emittance section
379
- ex = self.emittance[0]
380
- ey = self.emittance[1]
477
+ ex = self.nemitt_x
478
+ ey = self.nemitt_y
381
479
  file.write(f'\nemittance:\n x: {ex}\n y: {ey}\n')
382
480
 
383
481
  # Collimators section
384
482
  file.write('\ncollimators:\n')
385
483
  b1_colls, b2_colls, bx_colls = [], [], []
386
- for coll in self._colldb.index:
387
- if coll == 'tclia.4r2' or coll == 'tclia.4l8':
484
+ for coll in self.to_pandas().index:
485
+ if coll == 'tclia.4r2' or coll == 'tclia.4l8': # TODO: hardcoded!!!
388
486
  b1_colls.append(coll)
389
487
  b2_colls.append(coll)
390
488
  elif coll[-2:] == 'b1':
@@ -393,7 +491,7 @@ class CollimatorDatabase:
393
491
  b2_colls.append(coll)
394
492
  else:
395
493
  bx_colls.append(coll)
396
-
494
+
397
495
  # Handle special cases for collimators
398
496
  if (('tclia.4r2' in b1_colls) or ('tclia.4l8' in b1_colls)) and (len(b1_colls) <= 2):
399
497
  b1_colls = []
@@ -409,702 +507,104 @@ class CollimatorDatabase:
409
507
  _print_colls(bx_colls, self, 'bx', file)
410
508
  print('WARNING -- some collimators could not be assigned to b1 or b2. Tracking might not work with those collimators. Please manually change the output file if necessary.')
411
509
 
412
-
413
- def __init__(self, **kwargs):
414
- # Get all arguments
415
- for var in self._init_vars:
416
- if var in self._init_var_defaults:
417
- kwargs.setdefault(var, self._init_var_defaults[var])
418
- elif var not in kwargs.keys():
419
- raise ValueError(f"CollimatorDatabase is missing required argument '{var}'!")
420
-
421
- self._optics = pd.DataFrame(columns=['x', 'px', 'y', 'py', 'betx', 'bety', 'alfx', 'alfy', 'dx', 'dy'])
422
- self._parse_dict(kwargs['collimator_dict'], kwargs['family_dict'],
423
- kwargs['beam'], kwargs['_yaml_merged'], kwargs.get('ignore_crystals', True))
424
- self.emittance = [kwargs['nemitt_x'], kwargs['nemitt_y']]
425
- self._beta_gamma_rel = None
426
-
427
-
428
- def __getitem__(self, name):
429
- return CollimatorSettings(name, self._colldb)
430
-
431
- def to_pandas(self):
432
- return pd.DataFrame({
433
- 's_center': self.s_center,
434
- 'gap': self.gap,
435
- 'jaw': self.jaw,
436
- 'beam_size': self.beam_size,
437
- 'aligned_to': self.align_to,
438
- 'angle': self.angle,
439
- 'material': self.material,
440
- # 'offset': self.offset,
441
- # 'tilt': self.tilt,
442
- 'stage': self.stage,
443
- 'active_length': self.active_length,
444
- 'collimator_type': self.collimator_type,
445
- }, index=self.name)
446
-
447
-
448
- def _parse_dict(self, coll, fam, beam=None, _yaml_merged=False, ignore_crystals=True):
449
-
450
- # We make all keys case-insensitive to avoid confusion between different conventions
451
- coll = _dict_keys_to_lower(coll)
452
- fam = _dict_keys_to_lower(fam)
453
-
454
- # The dictionary can be a CollimatorDatabase for a single beam (beam=None)
455
- # or for both beams (beam='b1' or beam='b2)
456
- coll = _get_coll_dct_by_beam(coll, beam)
457
-
458
- # Apply family settings
459
- crystals = []
460
- for thiscoll, settings in coll.items():
461
- settings = {k.lower(): v for k,v in settings.items()}
462
- if 'family' in settings.keys() and settings['family'] is not None:
463
- settings['family'] = settings['family'].lower()
464
- thisfam = settings['family']
465
- if thisfam not in fam.keys():
466
- raise ValueError(f"Collimator {thiscoll} depends on family {thisfam}, "
467
- + f"but the latter is not defined!")
468
-
469
- # Check if some family settings are overwritten for this collimator
470
- # Only do this check if we didn't do a YAML merge earlier (because then it
471
- # is already taken care of)
472
- if not _yaml_merged:
473
- overwritten_keys = [key.lower() for key in settings.keys() if key in fam[thisfam]]
474
- if len(overwritten_keys) > 0:
475
- settings['overwritten_keys'] = overwritten_keys
476
-
477
- # Load family settings, potentially overwriting settings for this collimator
478
- settings = {**fam[thisfam], **settings}
479
-
510
+ # ====================================
511
+ # ====== Installing collimators ======
512
+ # ====================================
513
+
514
+ def install_black_absorbers(self, line, names=None, *, verbose=False, need_apertures=True):
515
+ self.line = line
516
+ elements = []
517
+ if names is None:
518
+ names = self.collimator_names
519
+ for name in names:
520
+ if verbose: print(f"Installing {name:20} as BlackAbsorber")
521
+ el = BlackAbsorber(gap=self[name]['gap'], angle=self[name]['angle'],
522
+ length=self[name]['length'], side=self[name]['side'],
523
+ _tracking=False)
524
+
525
+ # Check that collimator is not installed as different type
526
+ # TODO: automatically replace collimator type and print warning
527
+ if isinstance(line[name], tuple(collimator_classes)):
528
+ raise ValueError(f"Trying to install {name} as {el.__class__.__name__},"
529
+ + f" but it is already installed as {line[name].__class__.__name__}!\n"
530
+ + f"Please reconstruct the line.")
531
+
532
+ # TODO: only allow Marker elements, no Drifts!!
533
+ # How to do this with importing a line for MAD-X or SixTrack...?
534
+ elif not isinstance(line[name], (xt.Marker, xt.Drift)):
535
+ raise ValueError(f"Trying to install {name} as {el.__class__.__name__},"
536
+ + f" but the line element to replace is not an xtrack.Marker "
537
+ + f"(or xtrack.Drift)!\nPlease check the name, or correct the "
538
+ + f"element.")
539
+ el.emittance = [self.nemitt_x, self.nemitt_y]
540
+ elements.append(el)
541
+ install_elements(line, names, elements, need_apertures=need_apertures)
542
+
543
+ def install_everest_collimators(self, line, names=None, *, verbose=False, need_apertures=True):
544
+ self.line = line
545
+ elements = []
546
+ if names is None:
547
+ names = self.collimator_names
548
+ for name in names:
549
+ mat = SixTrack_to_xcoll[self[name]['material']]
550
+ if self[name]['crystal'] is None:
551
+ if verbose: print(f"Installing {name:20} as EverestCollimator")
552
+ el = EverestCollimator(gap=self[name]['gap'], angle=self[name]['angle'],
553
+ length=self[name]['length'], side=self[name]['side'],
554
+ material=mat[0], _tracking=False)
480
555
  else:
481
- settings['family'] = None
482
- coll[thiscoll] = settings
483
-
484
- # Save list of crystals
485
- if 'crystal' in settings:
486
- if settings['crystal'] != 0.0:
487
- crystals += [thiscoll]
488
- else:
489
- settings['crystal'] = None
490
-
491
- # Remove crystals from colldb
492
- if ignore_crystals:
493
- for thiscoll in crystals:
494
- del coll[thiscoll]
495
-
496
- # Check that all collimators have gap settings
497
- if not np.all(['gap' in val.keys() or 'opening' in val.keys() for val in coll.values()]):
498
- raise ValueError("Ill-defined CollimatorDatabase: Not all collimators have a gap or "
499
- + "opening setting, (or the keys / structure of the dictionary is wrong)!")
500
-
501
- # Update collimators with default values for missing keys
502
- for collimator in coll.values():
503
- # Change all values to lower case
504
- for key, val in collimator.items():
505
- collimator[key] = val.lower() if isinstance(val, str) else val
506
- if 'length' in collimator.keys():
507
- collimator['active_length'] = collimator.pop('length')
508
- if 'gap' in collimator.keys():
509
- if collimator['gap'] is not None and collimator['gap'] > 900:
510
- collimator['gap'] = None
511
- if 'side' in collimator.keys() and collimator['side'] == 'left':
512
- collimator['gap_L'] = collimator.pop('gap')
513
- collimator['gap_R'] = None
514
- elif 'side' in collimator.keys() and collimator['side'] == 'right':
515
- collimator['gap_L'] = None
516
- collimator['gap_R'] = collimator.pop('gap')
517
- else:
518
- collimator['gap_L'] = collimator['gap']
519
- collimator['gap_R'] = collimator.pop('gap')
520
- if 'angle' in collimator.keys():
521
- collimator['angle_L'] = collimator['angle']
522
- collimator['angle_R'] = collimator.pop('angle')
523
- _initialise_None(collimator)
524
-
525
- self._collimator_dict = coll
526
- self._family_dict = fam
527
- self._colldb = pd.DataFrame(coll).transpose()
528
-
529
-
530
- @property
531
- def name(self):
532
- return self._colldb.index.values
533
-
534
-
535
- # TODO: - VALIDATION OF TYPES (e.g. material, stage, align, ..)
536
- # - IMPLEMENTATION OF TILT
537
- # - CRYSTAL PROPERTIES: only valid if crystal == True (add mask to _set_property?)
538
- # make second dataframe for crystals
539
- # - show as __repr__
540
-
541
- # The CollimatorDatabase class has the following fields (those marked
542
- # with an * are set automatically and cannot be overwritten):
543
- # - name
544
- # - gap
545
- # - jaw *
546
- # - beam_size *
547
- # - s_center *
548
- # - angle
549
- # - material
550
- # - offset
551
- # - tilt
552
- # - stage
553
- # - side
554
- # - active_length
555
- # - inactive_front
556
- # - inactive_back
557
- # - total_length *
558
- # - collimator_type *
559
- # - betx
560
- # - bety
561
- # - x
562
- # - px
563
- # - y
564
- # - py
565
- # - gamma_rel
566
- # - emit
567
-
568
- @property
569
- def angle(self):
570
- # angles = np.array([self._colldb.angle_L.values,self._colldb.angle_R.values])
571
- # return pd.Series([ L if L == R else [L,R] for L, R in angles.T ], index=self._colldb.index, dtype=object)
572
- return self._colldb['angle_L']
573
-
574
- @angle.setter
575
- def angle(self, angle):
576
- self._set_property_LR('angle', angle)
577
- self._compute_jaws()
578
-
579
- @property
580
- def material(self):
581
- return self._colldb['material']
582
-
583
- @material.setter
584
- def material(self, material):
585
- self._set_property('material', material)
586
-
587
- @property
588
- def offset(self):
589
- return self._colldb['offset']
590
-
591
- @offset.setter
592
- def offset(self, offset):
593
- self._set_property('offset', offset, single_default_allowed=True)
594
- self._compute_jaws()
595
-
596
- @property
597
- def tilt(self):
598
- tilts = np.array([self._colldb.tilt_L.values,self._colldb.tilt_R.values])
599
- return pd.Series([ L if L == R else [L,R] for L, R in tilts.T ], index=self._colldb.index, dtype=object)
600
-
601
- @tilt.setter
602
- def tilt(self, tilts):
603
- self._set_property_LR('tilt', tilts)
604
- self._compute_jaws()
605
-
606
- @property
607
- def stage(self):
608
- return self._colldb['stage']
609
-
610
- @stage.setter
611
- def stage(self, stage):
612
- self._set_property('stage', stage)
613
-
614
- @property
615
- def parking(self):
616
- return self._colldb['parking']
617
-
618
- @parking.setter
619
- def parking(self, parking):
620
- self._set_property('parking', parking, single_default_allowed=True)
621
- self._compute_jaws()
622
-
623
- @property
624
- def active(self):
625
- return self._colldb['active']
626
-
627
- @active.setter
628
- def active(self, active):
629
- self._set_property('active', active, single_default_allowed=True)
630
-
631
- # @property
632
- # def crystal(self):
633
- # return self._colldb['crystal']
634
-
635
- # @crystal.setter
636
- # def crystal(self, crystal):
637
- # self._set_property('crystal', crystal)
638
-
639
- # @property
640
- # def bend(self):
641
- # return self._colldb['bend']
642
-
643
- # @bend.setter
644
- # def bend(self, bend):
645
- # self._set_property('bend', bend)
646
-
647
- # @property
648
- # def xdim(self):
649
- # return self._colldb['xdim']
650
-
651
- # @xdim.setter
652
- # def xdim(self, xdim):
653
- # self._set_property('xdim', xdim)
654
-
655
- # @property
656
- # def ydim(self):
657
- # return self._colldb['ydim']
658
-
659
- # @ydim.setter
660
- # def ydim(self, ydim):
661
- # self._set_property('ydim', ydim)
662
-
663
- # @property
664
- # def miscut(self):
665
- # return self._colldb['miscut']
666
-
667
- # @miscut.setter
668
- # def miscut(self, miscut):
669
- # self._set_property('miscut', miscut)
670
-
671
- # @property
672
- # def thick(self):
673
- # return self._colldb['thick']
674
-
675
- # @thick.setter
676
- # def thick(self, thick):
677
- # self._set_property('thick', thick)
678
-
679
- @property
680
- def s_center(self):
681
- return self._colldb['s_center']
682
-
683
- @property
684
- def collimator_type(self):
685
- return self._colldb['collimator_type']
686
-
687
- @property
688
- def active_length(self):
689
- return self._colldb['active_length']
690
-
691
- @active_length.setter
692
- def active_length(self, length):
693
- self._set_property('active_length', length)
694
- self.align_to = {}
695
-
696
- @property
697
- def inactive_front(self):
698
- return self._colldb['inactive_front']
699
-
700
- @inactive_front.setter
701
- def inactive_front(self, length):
702
- self._set_property('inactive_front', length)
703
-
704
- @property
705
- def inactive_back(self):
706
- return self._colldb['inactive_back']
707
-
708
- @inactive_back.setter
709
- def inactive_back(self, length):
710
- self._set_property('inactive_back', length)
711
-
712
- @property
713
- def total_length(self):
714
- return self._colldb['active_length'] + self._colldb['inactive_front'] + self._colldb['inactive_back']
715
-
716
- @property
717
- def gap(self):
718
- gaps = np.array([self._colldb.gap_L.values,self._colldb.gap_R.values])
719
- return pd.Series([ L if L == R else [L,R] for L, R in gaps.T ], index=self._colldb.index, dtype=object)
720
-
721
- @gap.setter
722
- def gap(self, gaps):
723
- df = self._colldb
724
- correct_format = False
725
- # The variable gaps is a Series or a list
726
- if isinstance(gaps, pd.Series) or isinstance(gaps, list) or isinstance(gaps, np.ndarray):
727
- correct_format = True
728
- if len(gaps) != len(self.name):
729
- raise ValueError("The variable 'gaps' has a different length than the number "
730
- + "of collimators in the CollimatorDatabase. Use a dictionary instead.")
731
- # Some of the gaps are list (e.g. two different values for both gaps): loop over gaps as dict
732
- if any(hasattr(gap, '__iter__') for gap in gaps):
733
- gaps = dict(zip(self.name, gaps))
734
- # All gaps are single values: use pandas-style assignment
735
- else:
736
- # mask those that have an active side for the gap under consideration
737
- # and have a setting less than 900; the others are set to None
738
- mask_L = np.logical_and(df.side.isin(['both','left']), ~(gaps >= 900))
739
- mask_R = np.logical_and(df.side.isin(['both','right']), ~(gaps >= 900))
740
- df.loc[mask_L, 'gap_L'] = gaps[mask_L]
741
- df.loc[~mask_L, 'gap_L'] = None
742
- df.loc[mask_R, 'gap_R'] = gaps[mask_R]
743
- df.loc[~mask_R, 'gap_R'] = None
744
-
745
- # The variable gaps is a dictionary
746
- if isinstance(gaps, dict):
747
- correct_format = True
748
- for name, gap in gaps.items():
749
- if name not in self.name:
750
- raise ValueError(f"Collimator {name} not found in CollimatorDatabase!")
751
- side = df.side[name]
752
- if hasattr(gap, '__iter__'):
753
- if isinstance(gap, str):
754
- raise ValueError("The gap setting has to be a number!")
755
- elif len(gap) == 2:
756
- gap_L = gap[0]
757
- gap_R = gap[1]
758
- if side != 'both':
759
- if side == 'left' and gap_R is not None:
760
- print(f"Warning: collimator {name} is left-sided but a finite right gap is specified. "
761
- + "Verify that this is what you want.")
762
- elif side == 'right' and gap_L is not None:
763
- print(f"Warning: collimator {name} is right-sided but a finite left gap is specified. "
764
- + "Verify that this is what you want.")
765
- elif len(gap) == 1:
766
- gap_L = gap[0] if side in ['both','left'] else None
767
- gap_R = gap[0] if side in ['both','right'] else None
768
- else:
769
- raise ValueError("The gap setting must have one or two values (for the left and the right jaw)!")
770
- else:
771
- gap_L = gap if side in ['both','left'] else None
772
- gap_R = gap if side in ['both','right'] else None
773
- gap_L = None if (gap_L is not None and gap_L >= 900) else gap_L
774
- gap_R = None if (gap_R is not None and gap_R >= 900) else gap_R
775
- df.loc[name, 'gap_L'] = gap_L
776
- df.loc[name, 'gap_R'] = gap_R
777
-
778
- if not correct_format:
779
- raise ValueError("Variable 'gaps' needs to be a pandas Series, dict, numpy array, or list!")
780
-
781
- df.gap_L = df.gap_L.astype('object', copy=False)
782
- df.gap_R = df.gap_R.astype('object', copy=False)
783
- self._compute_jaws()
784
-
785
- @property
786
- def jaw(self):
787
- jaws = list(np.array([
788
- self._colldb.jaw_LU.values,
789
- self._colldb.jaw_RU.values,
790
- self._colldb.jaw_LD.values,
791
- self._colldb.jaw_RD.values
792
- ]).T)
793
- # Need special treatment if there are None's
794
- def flip(jaw):
795
- return None if jaw is None else -jaw
796
- for i, jaw in enumerate(jaws):
797
- # All 4 jaw points are the same
798
- if jaw[0] == flip(jaw[1]) == jaw[2] == flip(jaw[3]):
799
- jaws[i] = jaw[0]
800
- # Upstream and downstream jaws are the same
801
- # (all cases except angular alignment and/or tilt)
802
- elif jaw[0] == jaw[2] and jaw[1] == jaw[3]:
803
- jaws[i] = [ jaw[0], jaw[1] ]
804
- else:
805
- jaws[i] = [ [jaw[0],jaw[1]], [jaw[2],jaw[3]] ]
806
- return pd.Series(jaws, index=self._colldb.index, dtype=object)
807
-
808
- @property
809
- def side(self):
810
- return self._colldb.side
811
-
812
- @side.setter
813
- def side(self, sides):
814
- self._set_property('side', sides, single_default_allowed=True)
815
- self.gap = self.gap
816
-
817
- @property
818
- def gamma_rel(self):
819
- return np.sqrt(self._beta_gamma_rel**2+1)
820
-
821
- @gamma_rel.setter
822
- def gamma_rel(self, gamma_rel):
823
- self._beta_gamma_rel = np.sqrt(gamma_rel**2-1)
824
- self._compute_jaws()
825
-
826
- @property
827
- def emittance(self):
828
- return [self._emitx, self._emity]
829
-
830
- @emittance.setter
831
- def emittance(self, emit):
832
- if hasattr(emit, '__iter__'):
833
- if isinstance(emit, str):
834
- raise ValueError(f"The 'emit' setting has to be a number!")
835
- elif len(emit) == 2:
836
- self._emitx = emit[0]
837
- self._emity = emit[1]
838
- elif len(emit) == 1:
839
- self._emitx = emit[0]
840
- self._emity = emit[0]
841
- else:
842
- raise ValueError(f"The 'emit' setting must have one or two values (for emitx and emity)!")
843
- else:
844
- self._emitx = emit
845
- self._emity = emit
846
- self._compute_jaws()
847
-
848
- @property
849
- def align_to(self):
850
- return self._colldb.align_to
851
-
852
- @align_to.setter
853
- def align_to(self, align):
854
- self._set_property('align_to', align, single_default_allowed=True, limit_to=['front', 'center', 'back', 'angular'])
855
- if np.any(self.align_to == 'maximum'):
856
- raise NotImplementedError
857
- s_front = self.s_center - self.active_length/2
858
- s_center = self.s_center
859
- s_back = self.s_center + self.active_length/2
860
- mask = self.align_to == 'front'
861
- self._colldb.loc[mask,'s_align_front'] = s_front[mask]
862
- self._colldb.loc[mask,'s_align_back'] = s_front[mask]
863
- mask = self.align_to == 'center'
864
- self._colldb.loc[mask,'s_align_front'] = s_center[mask]
865
- self._colldb.loc[mask,'s_align_back'] = s_center[mask]
866
- mask = self.align_to == 'back'
867
- self._colldb.loc[mask,'s_align_front'] = s_back[mask]
868
- self._colldb.loc[mask,'s_align_back'] = s_back[mask]
869
- mask = self.align_to == 'angular'
870
- self._colldb.loc[mask,'s_align_front'] = s_front[mask]
871
- self._colldb.loc[mask,'s_align_back'] = s_back[mask]
872
- self._compute_jaws()
873
-
874
- # TODO: when does this need to be unset?
875
- @property
876
- def _optics_is_ready(self):
877
- pos = set(self._colldb.s_align_front.values) | set(self._colldb.s_align_back.values)
878
- return np.all([s in self._optics.index for s in pos]) and self._beta_gamma_rel is not None
879
-
880
- @property
881
- def betx(self):
882
- vals = np.array([
883
- [ self._optics.loc[s,'betx'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
884
- [ self._optics.loc[s,'betx'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
885
- ])
886
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
887
-
888
- @property
889
- def bety(self):
890
- vals = np.array([
891
- [ self._optics.loc[s,'bety'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
892
- [ self._optics.loc[s,'bety'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
893
- ])
894
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
895
-
896
- @property
897
- def alfx(self):
898
- vals = np.array([
899
- [ self._optics.loc[s,'alfx'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
900
- [ self._optics.loc[s,'alfx'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
901
- ])
902
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
903
-
904
- @property
905
- def alfy(self):
906
- vals = np.array([
907
- [ self._optics.loc[s,'alfy'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
908
- [ self._optics.loc[s,'alfy'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
909
- ])
910
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
911
-
912
- @property
913
- def dx(self):
914
- vals = np.array([
915
- [ self._optics.loc[s,'dx'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
916
- [ self._optics.loc[s,'dx'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
917
- ])
918
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
919
-
920
- @property
921
- def dy(self):
922
- vals = np.array([
923
- [ self._optics.loc[s,'dy'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
924
- [ self._optics.loc[s,'dy'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
925
- ])
926
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
927
-
928
- @property
929
- def x(self):
930
- vals = np.array([
931
- [ self._optics.loc[s,'x'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
932
- [ self._optics.loc[s,'x'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
933
- ])
934
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
935
-
936
- @property
937
- def px(self):
938
- vals = np.array([
939
- [ self._optics.loc[s,'px'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
940
- [ self._optics.loc[s,'px'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
941
- ])
942
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
943
-
944
- @property
945
- def y(self):
946
- vals = np.array([
947
- [ self._optics.loc[s,'y'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
948
- [ self._optics.loc[s,'y'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
949
- ])
950
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
951
-
952
- @property
953
- def py(self):
954
- vals = np.array([
955
- [ self._optics.loc[s,'py'] if s in self._optics.index else None for s in self._colldb.s_align_front.values ],
956
- [ self._optics.loc[s,'py'] if s in self._optics.index else None for s in self._colldb.s_align_back.values ]
957
- ])
958
- return pd.Series([ F if F == B else [F,B] for F,B in vals.T ], index=self._colldb.index, dtype=object)
959
-
960
- @property
961
- def beam_size(self):
962
- if self._optics_is_ready:
963
- beam_size = np.array([self._beam_size_front,self._beam_size_back])
964
- return pd.Series([ F if F == B else [F,B] for F, B in beam_size.T ], index=self._colldb.index, dtype=object)
556
+ if verbose: print(f"Installing {name:20} as EverestCrystal")
557
+ el = EverestCrystal(gap=self[name]['gap'], angle=self[name]['angle'],
558
+ length=self[name]['length'], side=self[name]['side'], material=mat[1],
559
+ lattice=self[name]['crystal'], bending_radius=self[name]['bending_radius'],
560
+ width=self[name]['width'], height=self[name]['height'],
561
+ miscut=self[name]['miscut'], _tracking=False)
562
+
563
+ # Check that collimator is not installed as different type
564
+ # TODO: automatically replace collimator type and print warning
565
+ if isinstance(line[name], tuple(collimator_classes)):
566
+ raise ValueError(f"Trying to install {name} as {el.__class__.__name__},"
567
+ + f" but it is already installed as {line[name].__class__.__name__}!\n"
568
+ + f"Please reconstruct the line.")
569
+
570
+ # TODO: only allow Marker elements, no Drifts!!
571
+ # How to do this with importing a line for MAD-X or SixTrack...?
572
+ elif not isinstance(line[name], (xt.Marker, xt.Drift)):
573
+ raise ValueError(f"Trying to install {name} as {el.__class__.__name__},"
574
+ + f" but the line element to replace is not an xtrack.Marker "
575
+ + f"(or xtrack.Drift)!\nPlease check the name, or correct the "
576
+ + f"element.")
577
+ el.emittance = [self.nemitt_x, self.nemitt_y]
578
+ elements.append(el)
579
+ install_elements(line, names, elements, need_apertures=need_apertures)
580
+
581
+ # ==================================
582
+ # ====== Accessing attributes ======
583
+ # ==================================
584
+
585
+ @property
586
+ def collimator_names(self):
587
+ return list(self._collimator_dict.keys())
588
+
589
+ @property
590
+ def collimator_families(self):
591
+ return self._family_dict.keys()
592
+
593
+ def __getattr__(self, name):
594
+ if name in self._family_dict:
595
+ return self._family_dict[name]
596
+ elif name in self._collimator_dict:
597
+ return self._collimator_dict[name]
965
598
  else:
966
- return None
967
-
968
- @property
969
- def _beam_size_front(self):
970
- # TODO: curretnly only for angle_L
971
- df = self._colldb
972
- opt = self._optics
973
- betx = opt.loc[df.s_align_front,'betx'].astype(float)
974
- bety = opt.loc[df.s_align_front,'bety'].astype(float)
975
- sigmax = np.sqrt(betx*self._emitx/self._beta_gamma_rel)
976
- sigmay = np.sqrt(bety*self._emity/self._beta_gamma_rel)
977
- result = np.sqrt(
978
- (sigmax*np.cos(np.float_(df.angle_L.values)*np.pi/180))**2
979
- + (sigmay*np.sin(np.float_(df.angle_L.values)*np.pi/180))**2
980
- )
981
- result.index = self._colldb.index
982
- return result
599
+ raise ValueError(f"Family nor collimator '{name}' found in CollimatorDatabase!")
983
600
 
984
- @property
985
- def _beam_size_back(self):
986
- # TODO: curretnly only for angle_L
987
- df = self._colldb
988
- opt = self._optics
989
- betx = opt.loc[df.s_align_back,'betx'].astype(float)
990
- bety = opt.loc[df.s_align_back,'bety'].astype(float)
991
- sigmax = np.sqrt(betx*self._emitx/self._beta_gamma_rel)
992
- sigmay = np.sqrt(bety*self._emity/self._beta_gamma_rel)
993
- result = np.sqrt(
994
- (sigmax*np.cos(np.float_(df.angle_L.values)*np.pi/180))**2
995
- + (sigmay*np.sin(np.float_(df.angle_L.values)*np.pi/180))**2
996
- )
997
- result.index = self._colldb.index
998
- return result
999
-
1000
- # parking is defined with respect to closed orbit
1001
- # TODO: tilt
1002
- # 'upstr' => 'front' en 'downstr' => 'back'
1003
- def _compute_jaws(self):
1004
- if self._optics_is_ready:
1005
- df = self._colldb
1006
- beam_size_front = self._beam_size_front
1007
- beam_size_back = self._beam_size_back
1008
- jaw_LU = df['gap_L']*beam_size_front + self.offset
1009
- jaw_RU = df['gap_R']*beam_size_front - self.offset
1010
- jaw_LD = df['gap_L']*beam_size_back + self.offset
1011
- jaw_RD = df['gap_R']*beam_size_back - self.offset
1012
- df['jaw_LU'] = df['parking'] if df['gap_L'] is None else np.minimum(jaw_LU,df['parking'])
1013
- df['jaw_RU'] = -df['parking'] if df['gap_R'] is None else -np.minimum(jaw_RU,df['parking'])
1014
- df['jaw_LD'] = df['parking'] if df['gap_L'] is None else np.minimum(jaw_LD,df['parking'])
1015
- df['jaw_RD'] = -df['parking'] if df['gap_R'] is None else -np.minimum(jaw_RD,df['parking'])
1016
- # align crystals
1017
- opt = self._optics
1018
- df['align_angle'] = None
1019
- cry_mask = [c is not None for c in df.crystal]
1020
- df_cry = df[cry_mask]
1021
- if len(df_cry) > 0:
1022
- alfx = opt.loc[df_cry.s_align_front,'alfx'].astype(float).values
1023
- alfy = opt.loc[df_cry.s_align_front,'alfy'].astype(float).values
1024
- betx = opt.loc[df_cry.s_align_front,'betx'].astype(float).values
1025
- bety = opt.loc[df_cry.s_align_front,'bety'].astype(float).values
1026
- align_angle_x = -np.sqrt(self._emitx/self._beta_gamma_rel/betx)*alfx
1027
- align_angle_y = -np.sqrt(self._emity/self._beta_gamma_rel/bety)*alfy
1028
- align_angle = np.array([x if abs(ang) < 1e-6 else y
1029
- for x,y,ang in zip(align_angle_x,align_angle_y,df_cry.angle_L.values)])
1030
- df.loc[cry_mask, 'align_angle'] = align_angle*df_cry['gap_L']
1031
-
1032
-
1033
- # ---------------------------------------
1034
- # ------ Property setter functions ------
1035
- # ---------------------------------------
1036
-
1037
- def _set_property(self, prop, vals, single_default_allowed=False, limit_to=[]):
1038
- df = self._colldb
1039
- if not isinstance(limit_to, (list, tuple, set)):
1040
- limit_to = [limit_to]
1041
- if isinstance(vals, dict):
1042
- for name, val in vals.items():
1043
- if name not in self.name:
1044
- raise ValueError(f"Collimator {name} not found in CollimatorDatabase!")
1045
- if limit_to!=[] and val not in limit_to:
1046
- raise ValueError(f"Cannot set {prop} to {val}. Choose from {limit_to}!")
1047
- df.loc[name, prop] = val
1048
- elif isinstance(vals, pd.Series) or isinstance(vals, list) or isinstance(vals, np.ndarray):
1049
- if len(vals) != len(self.name):
1050
- raise ValueError(f"The variable '{prop}' has a different length than the number of "
1051
- + "collimators in the CollimatorDatabase. Use a dictionary instead.")
1052
- if limit_to!=[] and np.any([val not in limit_to for val in vals]):
1053
- raise ValueError(f"Cannot set {prop} to {vals}. Choose from {limit_to}!")
1054
- df[prop] = vals
601
+ def __getitem__(self, name):
602
+ if name in self._family_dict:
603
+ return self._family_dict[name]
604
+ elif name in self._collimator_dict:
605
+ return self._collimator_dict[name]
1055
606
  else:
1056
- if single_default_allowed:
1057
- if limit_to!=[] and vals not in limit_to:
1058
- raise ValueError(f"Cannot set {prop} to {vals}. Choose from {limit_to}!")
1059
- df[prop] = vals
1060
- else:
1061
- raise ValueError(f"Variable '{prop}' needs to be a pandas Series, dict, numpy array, or list!")
1062
-
1063
-
1064
- def _set_property_LR(self, prop, vals):
1065
- df = self._colldb
1066
- correct_format = False
1067
- # The variable vals is a Series or a list
1068
- if isinstance(vals, pd.Series) or isinstance(vals, list) or isinstance(vals, np.ndarray):
1069
- correct_format = True
1070
- if len(vals) != len(self.name):
1071
- raise ValueError(f"The variable '{prop}' has a different length than the number of "
1072
- + "collimators in the CollimatorDatabase. Use a dictionary instead.")
1073
- # Some of the vals are list (e.g. two different values for both gaps): loop over vals as dict
1074
- if any(hasattr(val, '__iter__') for val in vals):
1075
- vals = dict(zip(self.name, vals))
1076
- # All gaps are single values: use pandas-style assignment
1077
- else:
1078
- df[prop + "_L"] = vals
1079
- df[prop + "_R"] = vals
607
+ raise ValueError(f"Family nor collimator '{name}' found in CollimatorDatabase!")
1080
608
 
1081
- # The variable vals is a dictionary
1082
- if isinstance(vals, dict):
1083
- correct_format = True
1084
- for name, val in vals.items():
1085
- if name not in self.name:
1086
- raise ValueError(f"Collimator {name} not found in CollimatorDatabase!")
1087
- if hasattr(val, '__iter__'):
1088
- if isinstance(val, str):
1089
- raise ValueError(f"The '{prop}' setting has to be a number!")
1090
- elif len(val) == 2:
1091
- val_L = val[0]
1092
- val_R = val[1]
1093
- elif len(val) == 1:
1094
- val_L = val[0]
1095
- val_R = val[0]
1096
- else:
1097
- raise ValueError(f"The '{prop}' setting must have one or two values (for the left and the right jaw)!")
1098
- else:
1099
- val_L = val
1100
- val_R = val
1101
- df.loc[name, prop + "_L"] = val_L
1102
- df.loc[name, prop + "_R"] = val_R
1103
-
1104
- if not correct_format:
1105
- raise ValueError("Variable '{prop}' needs to be a pandas Series, dict, numpy array, or list!")
1106
-
1107
- df[prop + "_L"] = df[prop + "_L"].astype('object', copy=False)
1108
- df[prop + "_R"] = df[prop + "_R"].astype('object', copy=False)
1109
- # Check if collimator active
1110
- # Check if gap is list (assymetric jaws)
609
+
610
+