Skip to content

Commit

Permalink
Merge branch 'fix-kernel-order' into release/v0.54.1
Browse files Browse the repository at this point in the history
  • Loading branch information
@szymonlopaciuk
szymonlopaciuk committed Mar 4, 2024
2 parents 90d2ab4 + a5f49c8 commit b532415
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 86 deletions.
49 changes: 32 additions & 17 deletions tests/test_prebuild_kernels.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@
from xtrack.prebuild_kernels import regenerate_kernels


def test_prebuild_kernels(mocker, tmp_path, temp_context_default_func):
def test_prebuild_kernels(mocker, tmp_path, temp_context_default_func, capsys):
# Set up the temporary kernels directory
kernel_definitions = {
"test_module": {
kernel_definitions = [
("111_test_module", {
"config": {
"XTRACK_MULTIPOLE_NO_SYNRAD": True,
"XTRACK_GLOBAL_XY_LIMIT": 1.0,
Expand All @@ -25,9 +25,23 @@ def test_prebuild_kernels(mocker, tmp_path, temp_context_default_func):
xt.Drift,
xt.Cavity,
xt.XYShift,
]
},
}
],
}),
("000_test_module", {
"config": {
"XTRACK_MULTIPOLE_NO_SYNRAD": True,
"XTRACK_GLOBAL_XY_LIMIT": 1.0,
"XFIELDS_BB3D_NO_BEAMSTR": True,
"XFIELDS_BB3D_NO_BHABHA": True,
},
"classes": [
xt.Drift,
xt.Cavity,
xt.XYShift,
xt.Bend,
],
}),
]

patch_defs = 'xtrack.prebuilt_kernels.kernel_definitions.kernel_definitions'
mocker.patch(patch_defs, kernel_definitions)
Expand All @@ -40,19 +54,17 @@ def test_prebuild_kernels(mocker, tmp_path, temp_context_default_func):
regenerate_kernels()

# Check if the expected files were created
so_file_exists = False
for path in tmp_path.iterdir():
if not path.name.startswith('test_module.'):
continue
if path.suffix not in ('.so', '.dll', '.dylib', '.pyd'):
continue
so_file_exists = True
assert so_file_exists
so_file0, = tmp_path.glob('000_test_module.*.so')
assert so_file0.exists()
assert (tmp_path / '000_test_module.c').exists()
assert (tmp_path / '000_test_module.json').exists()

assert (tmp_path / 'test_module.c').exists()
assert (tmp_path / 'test_module.json').exists()
so_file1, = tmp_path.glob('111_test_module.*.so')
assert so_file1.exists()
assert (tmp_path / '111_test_module.c').exists()
assert (tmp_path / '111_test_module.json').exists()

# Test that reloading the kernel works
# Test that reloading the kernel works, and it's the right one
cffi_compile = mocker.patch.object(cffi.FFI, 'compile')

line = xt.Line(elements=[xt.Drift(length=2.0)])
Expand All @@ -67,6 +79,9 @@ def test_prebuild_kernels(mocker, tmp_path, temp_context_default_func):
assert p.y == 0.0
cffi_compile.assert_not_called()

captured = capsys.readouterr()
assert 'Found suitable prebuilt kernel `111_test_module`' in captured.out


def test_per_element_prebuild_kernels(mocker, tmp_path, temp_context_default_func):
# Set up the temporary kernels directory
Expand Down
9 changes: 6 additions & 3 deletions xtrack/prebuild_kernels.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,8 +143,11 @@ def enumerate_kernels() -> Iterator[Tuple[str, dict]]:
xsuite. The first element of the tuple is the name of the kernel module
and the second is a dictionary with the kernel metadata.
"""
for metadata_file in XT_PREBUILT_KERNELS_LOCATION.glob('*.json'):
if metadata_file.stem.startswith('_'):
from xtrack.prebuilt_kernels.kernel_definitions import kernel_definitions
for kernel_name, _ in kernel_definitions:
metadata_file = XT_PREBUILT_KERNELS_LOCATION / f'{kernel_name}.json'

if not metadata_file.exists():
continue

with metadata_file.open('r') as fd:
Expand Down Expand Up @@ -277,7 +280,7 @@ def regenerate_kernels(kernels=None):
except ImportError:
pass

for module_name, metadata in kernel_definitions.items():
for module_name, metadata in kernel_definitions:
if kernels is not None and module_name not in kernels:
continue

Expand Down
144 changes: 78 additions & 66 deletions xtrack/prebuilt_kernels/kernel_definitions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,87 +68,94 @@
MultiSetter
]

# These will be enumerated in order of appearance in the dict, so in this case
# (for optimization purposes) the order is important.
kernel_definitions = {
'default_only_xtrack_no_config': {
# These are enumerated in order specified below: the highest priority at the top
kernel_definitions = [
('default_only_xtrack_no_config', {
'config': {},
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'default_only_xtrack': {
}),
('default_only_xtrack', {
'config': BASE_CONFIG,
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'only_xtrack_non_tracking_kernels': {
'config': {},
}),
('only_xtrack_non_tracking_kernels', {
'config': BASE_CONFIG,
'classes': [],
'extra_classes': NON_TRACKING_ELEMENTS
},
'default_only_xtrack_backtrack': {
}),
('default_only_xtrack_backtrack', {
'config': {**BASE_CONFIG, 'XSUITE_BACKTRACK': True},
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'default_only_xtrack_backtrack_no_limit': {
'config': {**{k: v for k,v in BASE_CONFIG.items()
if k != 'XTRACK_GLOBAL_XY_LIMIT'},
'XSUITE_BACKTRACK': True},
}),
('default_only_xtrack_backtrack_no_limit', {
'config': {
**{k: v for k, v in BASE_CONFIG.items()
if k != 'XTRACK_GLOBAL_XY_LIMIT'},
'XSUITE_BACKTRACK': True
},
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'only_xtrack_frozen_longitudinal': {
}),
('only_xtrack_frozen_longitudinal', {
'config': {**BASE_CONFIG, **FREEZE_LONGITUDINAL},
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'only_xtrack_frozen_energy': {
}),
('only_xtrack_frozen_energy', {
'config': {**BASE_CONFIG, **FREEZE_ENERGY},
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'only_xtrack_backtrack_frozen_energy': {
}),
('only_xtrack_backtrack_frozen_energy', {
'config': {**BASE_CONFIG, **FREEZE_ENERGY, 'XSUITE_BACKTRACK': True},
'classes': ONLY_XTRACK_ELEMENTS + NO_SYNRAD_ELEMENTS,
},
'only_xtrack_taper': {
}),
('only_xtrack_taper', {
'config': {
**BASE_CONFIG,
'XTRACK_MULTIPOLE_NO_SYNRAD': False,
'XTRACK_MULTIPOLE_TAPER': True,
'XTRACK_DIPOLEEDGE_TAPER': True,
},
'classes': ONLY_XTRACK_ELEMENTS,
},
'only_xtrack_with_synrad': {
}),
('only_xtrack_with_synrad', {
'config': {**BASE_CONFIG, 'XTRACK_MULTIPOLE_NO_SYNRAD': False},
'classes': ONLY_XTRACK_ELEMENTS,
},
'only_xtrack_with_synrad_kick_as_co': {
'config': {**BASE_CONFIG, 'XTRACK_MULTIPOLE_NO_SYNRAD': False,
'XTRACK_SYNRAD_KICK_SAME_AS_FIRST': True},
}),
('only_xtrack_with_synrad_kick_as_co', {
'config': {
**BASE_CONFIG, 'XTRACK_MULTIPOLE_NO_SYNRAD': False,
'XTRACK_SYNRAD_KICK_SAME_AS_FIRST': True
},
'classes': ONLY_XTRACK_ELEMENTS,
}
}
}),
]


try:
import xfields as xf

DEFAULT_BB3D_ELEMENTS = [
*ONLY_XTRACK_ELEMENTS,
xf.BeamBeamBiGaussian2D,
xf.BeamBeamBiGaussian3D,
]

kernel_definitions['default_bb3d'] = {
kernel_definitions.append(('default_bb3d', {
'config': BASE_CONFIG,
'classes': [*DEFAULT_BB3D_ELEMENTS, LineSegmentMap],
}
kernel_definitions['default_bb3d_no_config'] = {
}))

kernel_definitions.append(('default_bb3d_no_config', {
'config': {},
'classes': [*DEFAULT_BB3D_ELEMENTS, LineSegmentMap],
}
}))

except ImportError:
LOGGER.warning('Xfields not installed, skipping BB3D elements')


try:
import xcoll as xc

DEFAULT_XCOLL_ELEMENTS = [
*ONLY_XTRACK_ELEMENTS,
*NO_SYNRAD_ELEMENTS,
Expand All @@ -158,36 +165,41 @@
xc.EverestCrystal
]

kernel_definitions['default_xcoll'] = {
'config': BASE_CONFIG,
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions['default_xcoll_no_config'] = {
'config': {},
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions['default_xcoll_frozen_longitudinal'] = {
'config': {**BASE_CONFIG, **FREEZE_LONGITUDINAL},
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions['default_xcoll_frozen_energy'] = {
'config': {**BASE_CONFIG, **FREEZE_ENERGY},
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions['default_xcoll_backtrack'] = {
'config': {**BASE_CONFIG, 'XSUITE_BACKTRACK': True},
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions['default_xcoll_backtrack_no_limit'] = {
'config': {**{k: v for k,v in BASE_CONFIG.items()
if k != 'XTRACK_GLOBAL_XY_LIMIT'},
'XSUITE_BACKTRACK': True},
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions['default_xcoll_backtrack_frozen_energy'] = {
'config': {**BASE_CONFIG, **FREEZE_ENERGY, 'XSUITE_BACKTRACK': True},
'classes': DEFAULT_XCOLL_ELEMENTS
}
kernel_definitions += [
('default_xcoll', {
'config': BASE_CONFIG,
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
('default_xcoll_no_config', {
'config': {},
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
('default_xcoll_frozen_longitudinal', {
'config': {**BASE_CONFIG, **FREEZE_LONGITUDINAL},
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
('default_xcoll_frozen_energy', {
'config': {**BASE_CONFIG, **FREEZE_ENERGY},
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
('default_xcoll_backtrack', {
'config': {**BASE_CONFIG, 'XSUITE_BACKTRACK': True},
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
('default_xcoll_backtrack_no_limit', {
'config': {
**{k: v for k, v in BASE_CONFIG.items()
if k != 'XTRACK_GLOBAL_XY_LIMIT'},
'XSUITE_BACKTRACK': True
},
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
('default_xcoll_backtrack_frozen_energy', {
'config': {**BASE_CONFIG, **FREEZE_ENERGY, 'XSUITE_BACKTRACK': True},
'classes': DEFAULT_XCOLL_ELEMENTS,
}),
]

except ImportError:
LOGGER.warning('Xcoll not installed, skipping collimator elements')

0 comments on commit b532415

Please sign in to comment.