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.
- xcoll/__init__.py +12 -4
- xcoll/beam_elements/__init__.py +7 -5
- xcoll/beam_elements/absorber.py +41 -7
- xcoll/beam_elements/base.py +1161 -244
- xcoll/beam_elements/collimators_src/black_absorber.h +118 -0
- xcoll/beam_elements/collimators_src/black_crystal.h +111 -0
- xcoll/beam_elements/collimators_src/everest_block.h +40 -28
- xcoll/beam_elements/collimators_src/everest_collimator.h +129 -50
- xcoll/beam_elements/collimators_src/everest_crystal.h +217 -73
- xcoll/beam_elements/everest.py +60 -113
- xcoll/colldb.py +250 -750
- xcoll/general.py +2 -2
- xcoll/headers/checks.h +1 -1
- xcoll/headers/particle_states.h +2 -2
- xcoll/initial_distribution.py +195 -0
- xcoll/install.py +177 -0
- xcoll/interaction_record/__init__.py +1 -0
- xcoll/interaction_record/interaction_record.py +252 -0
- xcoll/interaction_record/interaction_record_src/interaction_record.h +98 -0
- xcoll/{impacts → interaction_record}/interaction_types.py +11 -4
- xcoll/line_tools.py +83 -0
- xcoll/lossmap.py +209 -0
- xcoll/manager.py +2 -937
- xcoll/rf_sweep.py +1 -1
- xcoll/scattering_routines/everest/amorphous.h +239 -0
- xcoll/scattering_routines/everest/channeling.h +245 -0
- xcoll/scattering_routines/everest/crystal_parameters.h +137 -0
- xcoll/scattering_routines/everest/everest.h +8 -30
- xcoll/scattering_routines/everest/everest.py +13 -10
- xcoll/scattering_routines/everest/jaw.h +27 -197
- xcoll/scattering_routines/everest/materials.py +2 -0
- xcoll/scattering_routines/everest/multiple_coulomb_scattering.h +31 -10
- xcoll/scattering_routines/everest/nuclear_interaction.h +86 -0
- xcoll/scattering_routines/geometry/__init__.py +6 -0
- xcoll/scattering_routines/geometry/collimator_geometry.h +219 -0
- xcoll/scattering_routines/geometry/crystal_geometry.h +150 -0
- xcoll/scattering_routines/geometry/geometry.py +26 -0
- xcoll/scattering_routines/geometry/get_s.h +92 -0
- xcoll/scattering_routines/geometry/methods.h +111 -0
- xcoll/scattering_routines/geometry/objects.h +154 -0
- xcoll/scattering_routines/geometry/rotation.h +23 -0
- xcoll/scattering_routines/geometry/segments.h +226 -0
- xcoll/scattering_routines/geometry/sort.h +184 -0
- {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/METADATA +1 -1
- {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/RECORD +48 -33
- xcoll/beam_elements/collimators_src/absorber.h +0 -141
- xcoll/collimator_settings.py +0 -457
- xcoll/impacts/__init__.py +0 -1
- xcoll/impacts/impacts.py +0 -102
- xcoll/impacts/impacts_src/impacts.h +0 -99
- xcoll/scattering_routines/everest/crystal.h +0 -1302
- xcoll/scattering_routines/everest/scatter.h +0 -169
- xcoll/scattering_routines/everest/scatter_crystal.h +0 -260
- {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/LICENSE +0 -0
- {xcoll-0.3.6.dist-info → xcoll-0.4.0.dist-info}/NOTICE +0 -0
- {xcoll-0.3.6.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,
|
|
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
|
-
|
|
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
|
-
|
|
21
|
-
fields
|
|
22
|
-
fields.update({'
|
|
23
|
-
fields.update({'
|
|
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
|
|
29
|
-
|
|
30
|
-
for key in
|
|
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
|
-
|
|
71
|
-
|
|
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', '
|
|
269
|
+
cry_fields = ['bending_radius', 'width', 'height', 'thick', 'miscut', 'crystal']
|
|
197
270
|
cry = {}
|
|
198
271
|
|
|
199
|
-
for
|
|
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':
|
|
229
|
-
|
|
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:
|
|
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
|
-
|
|
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
|
|
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.
|
|
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','
|
|
309
|
-
if coll_dict['angle_L'] == coll_dict['angle_R']:
|
|
310
|
-
|
|
311
|
-
else:
|
|
312
|
-
|
|
313
|
-
if coll_dict['gap_L'] == coll_dict['gap_R']:
|
|
314
|
-
|
|
315
|
-
elif coll_dict['gap_L'] is None and coll_dict['gap_R'] is not None:
|
|
316
|
-
|
|
317
|
-
elif coll_dict['gap_L'] is not None and coll_dict['gap_R'] is None:
|
|
318
|
-
|
|
319
|
-
else:
|
|
320
|
-
|
|
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.
|
|
380
|
-
ey = self.
|
|
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.
|
|
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
|
-
|
|
414
|
-
|
|
415
|
-
|
|
416
|
-
|
|
417
|
-
|
|
418
|
-
|
|
419
|
-
|
|
420
|
-
|
|
421
|
-
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
433
|
-
|
|
434
|
-
|
|
435
|
-
|
|
436
|
-
|
|
437
|
-
|
|
438
|
-
|
|
439
|
-
|
|
440
|
-
|
|
441
|
-
|
|
442
|
-
|
|
443
|
-
|
|
444
|
-
|
|
445
|
-
|
|
446
|
-
|
|
447
|
-
|
|
448
|
-
|
|
449
|
-
|
|
450
|
-
|
|
451
|
-
|
|
452
|
-
|
|
453
|
-
|
|
454
|
-
|
|
455
|
-
|
|
456
|
-
|
|
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
|
-
|
|
482
|
-
|
|
483
|
-
|
|
484
|
-
|
|
485
|
-
|
|
486
|
-
|
|
487
|
-
|
|
488
|
-
|
|
489
|
-
|
|
490
|
-
|
|
491
|
-
|
|
492
|
-
|
|
493
|
-
|
|
494
|
-
|
|
495
|
-
|
|
496
|
-
|
|
497
|
-
|
|
498
|
-
|
|
499
|
-
|
|
500
|
-
|
|
501
|
-
|
|
502
|
-
|
|
503
|
-
|
|
504
|
-
|
|
505
|
-
|
|
506
|
-
|
|
507
|
-
|
|
508
|
-
|
|
509
|
-
|
|
510
|
-
|
|
511
|
-
|
|
512
|
-
|
|
513
|
-
|
|
514
|
-
|
|
515
|
-
|
|
516
|
-
|
|
517
|
-
|
|
518
|
-
|
|
519
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
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
|
-
|
|
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
|
-
|
|
985
|
-
|
|
986
|
-
|
|
987
|
-
|
|
988
|
-
|
|
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
|
-
|
|
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
|
-
|
|
1082
|
-
|
|
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
|
+
|