def test_inconsistent_energy_values_warns_correctly():
d1 = elements.Dipole('d1', 1, BendingAngle=numpy.pi, Energy=5.e+9,
PassMethod='BndMPoleSymplectic4RadPass')
d2 = elements.Dipole('d2', 1, BendingAngle=numpy.pi, Energy=3.e+9,
PassMethod='BndMPoleSymplectic4RadPass')
m1 = elements.Marker('m1', Energy=5.e+9)
m2 = elements.Marker('m2', Energy=3.e+9)
with pytest.warns(AtWarning):
params = _matlab_scanner([m1, m2])
assert params['_energy'] == 5.e+9
with pytest.warns(AtWarning):
params = _matlab_scanner([d1, d2])
assert params['_energy'] == 5.e+9
def ringparam_filter(params, elem_iterator, *args):
""""
Run through all elements, process and optionally removes RingParam elements
KEYWORDS
keep_all=False if True, keep RingParam elem_iterator as Markers
"""
keep_all = params.pop('keep_all', False)
ringparams = []
radiate = False
for elem in elem_iterator(params, *args):
if (elem.PassMethod.endswith('RadPass')
or elem.PassMethod.endswith('CavityPass')):
radiate = True
if isinstance(elem, RingParam):
ringparams.append(elem)
for k, v in elem.items():
if k not in _param_ignore:
params.setdefault(_param_to_lattice.get(k, k.lower()), v)
if keep_all:
pars = vars(elem).copy()
name = pars.pop('FamName')
yield elements.Marker(name, **pars)
else:
yield elem
params['_radiation'] = radiate
if len(ringparams) > 1:
warn(AtWarning('More than 1 RingParam element, the 1st one is used'))
def sbreak(self, break_s, break_elems=None):
"""Insert elements at selected locations in the lattice
PARAMETERS
break_s: location or array of locations of breakpoints
break_elems: elements to be inserted at breakpoints (array of
elements as long as break_s or single element
duplicated as necessary). Default: Marker('sbreak')
RETURNS
A new lattice with new elements inserted at breakpoints
"""
def sbreak_iterator(elems, insertions):
"""Iterate over elements and breaks where necessary"""
def next_mk():
"""Extract the next element to insert"""
try:
return next(insertions)
except StopIteration:
return sys.float_info.max, None
s_end = 0.0
# get the 1st insertion
smk, mk = next_mk()
# skip all insertions at negative break_s, if any
while smk < s_end:
smk, mk = next_mk()
for elem in elems:
s_end += elem.Length
# loop over all insertions within the element
while smk < s_end:
frac = (s_end - smk) / elem.Length
if frac < 1.0: # breakpoint is within the element
el0, elem = elem.divide([1.0 - frac, frac])
yield el0
yield mk
smk, mk = next_mk()
yield elem
# set default insertion
if break_elems is None:
break_elems = elements.Marker('sbreak')
break_elems = numpy.reshape(break_elems, -1)
# Check element lengths
if not all(e.Length == 0 for e in break_elems):
warn(
AtWarning(
"Inserting elements with length!=0 may change the lattice")
)
# broadcast break_s and break_elems to arrays of same size
# and create an iterator over the elements to be inserted
iter_mk = zip(*numpy.broadcast_arrays(break_s, break_elems))
return Lattice(sbreak_iterator(self, iter_mk), **vars(self))
def test_get_elements(hmba_lattice):
# test FamName direct match
assert get_elements(hmba_lattice, 'BPM_06') == [hmba_lattice[65]]
# test FamName wildcard matching
assert get_elements(hmba_lattice, 'QD2?') == list(hmba_lattice[9, 113])
assert get_elements(hmba_lattice, 'QD3*') == list(hmba_lattice[19, 105])
assert get_elements(hmba_lattice, 'S*H2B') == list([hmba_lattice[55]])
assert get_elements(hmba_lattice, '*C_1') == list(hmba_lattice[59, 60])
assert get_elements(hmba_lattice,
'DR_2[1-3]') == list(hmba_lattice[54, 56, 58])
assert get_elements(hmba_lattice,
'DR_2[!1-7]') == list(hmba_lattice[52, 78, 80])
# test element instance
marker = elements.Marker('M1')
assert get_elements(hmba_lattice,
marker) == list(hmba_lattice[1, 12, 61, 67, 73])
# test element type
assert get_elements(hmba_lattice, elements.RFCavity) == [hmba_lattice[0]]
# test invalid key raises TypeError
with pytest.raises(TypeError):
get_elements(hmba_lattice, None)
# test quiet suppresses print statement correctly
if sys.version_info < (3, 0):
capturedOutput = BytesIO()
else:
capturedOutput = StringIO()
sys.stdout = capturedOutput
get_elements(hmba_lattice, 'BPM_06', quiet=True)
sys.stdout = sys.__stdout__
assert capturedOutput.getvalue() == ''
sys.stdout = capturedOutput
get_elements(hmba_lattice, 'BPM_06', quiet=False)
sys.stdout = sys.__stdout__
assert capturedOutput.getvalue() == ("String 'BPM_06' matched 1 family: "
"BPM_06\nall corresponding elements "
"have been returned.\n")