def test_all_args(self):
el0 = madseq.Element(None, None, odicti(a='a0', b='b0', c='c0'))
el1 = madseq.Element(None, None, odicti(a='a1', b='b1', d='d1'), el0)
el2 = madseq.Element(None, None, odicti(a='a2'), el1)
self.assertEqual(el2.all_args,
dicti([('c', 'c0'), ('b', 'b1'),
('d', 'd1'), ('a', 'a2')]))
def test_replace_with_parent(self):
init_l = Decimal('1.5')
base = madseq.Element('BASE', 'DRIFT', odicti(l=init_l, k=2))
elem = madseq.Element(None, 'BASE', odicti(), base)
transformer = madseq.ElementTransform({})
tpl, el, l = transformer.slice(elem, 0, 0)
self.assertEqual(l, init_l)
def test_get(self):
el0 = madseq.Element(None, None, odicti(a=1))
el1 = madseq.Element(None, None, odicti(b=2, c=3), el0)
self.assertEqual(el1.get('a', 9), 1)
self.assertEqual(el1.get('b', 9), 2)
self.assertEqual(el1.get('c', 9), 3)
self.assertEqual(el1.get('x', 9), 9)
self.assertEqual(el1.get('x'), None)
def test_getitem(self):
el0 = madseq.Element(None, None, odicti(a='a0', b='b0', c='c0'))
el1 = madseq.Element(None, None, odicti(a='a1', b='b1', d='d1'), el0)
el2 = madseq.Element(None, None, odicti(a='a2'), el1)
self.assertEqual(el2['a'], 'a2')
self.assertEqual(el2['b'], 'b1')
self.assertEqual(el2['c'], 'c0')
self.assertEqual(el2['d'], 'd1')
def test_delitem(self):
el0 = madseq.Element(None, None, odicti(a=1))
el1 = madseq.Element(None, None, odicti(b=2, c=3), el0)
def delitem(d, key):
del d[key]
self.assertRaises(KeyError, delitem, el1, 'a')
del el1['b']
self.assertEqual(el1.args, odicti(c=3))
def test_parse_args(self):
parse_args = madseq.parse_args
self.assertEqual(parse_args(''),
odicti())
self.assertEqual(parse_args(', a := 3'),
odicti([('a', 3)]))
self.assertEqual(parse_args(', a =3, crd := pi'),
odicti([('a', 3), ('crd', pi)]))
def _getstate(self):
"""Get a serializeable state for :class:`Json` and :class:`Yaml`."""
return odicti(
(seq.name, odicti(
list(seq.head.args.items()) +
[('elements', [elem._getstate()
for elem in seq.body
if elem.type])]
))
for seq in self._nodes
if isinstance(seq, Sequence))
def test_pop(self):
el0 = madseq.Element(None, None, odicti(a=1))
el1 = madseq.Element(None, None, odicti(b=2, c=3), el0)
self.assertRaises(KeyError, el1.pop, 'x')
self.assertEqual(el1.pop('x', 3), 3)
self.assertEqual(el1.pop('a'), 1)
self.assertEqual(el0['a'], 1) # el0 must not be modified!
self.assertEqual(el1.pop('b'), 2)
self.assertEqual(el1.args, odicti(c=3))
def test_copy(self):
el = madseq.Element(None, None, odicti(a=1))
copy = el.copy()
self.assertFalse(el is copy)
self.assertFalse(el.args is copy.args)
self.assertEqual(el, copy)
self.assertEqual(el.args, copy.args)
def all_args(self):
"""Return merged arguments of self and bases."""
if self._base:
args = self._base.all_args
else:
args = odicti()
args.update(self.args)
return args
def load(self, stream):
"""Load data from, using ordered case insensitive dictionaries."""
yaml = self.yaml
class OrderedLoader(yaml.SafeLoader):
pass
OrderedLoader.add_constructor(
yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
lambda loader, node: odicti(loader.construct_pairs(node)))
return yaml.load(stream, OrderedLoader)
def test_solenoid(self):
l = 3.5
pi = 3.14
ratio = 0.79
sol = madseq.Element(None, 'solenoid', odicti(Ks=pi, L=l))
scaled = madseq.rescale_makethin(sol, ratio)
self.assertEqual(scaled['ksi'], pi*l*ratio)
self.assertEqual(scaled['lrad'], l*ratio)
self.assertEqual(scaled.type, 'solenoid')
def test_quadrupole(self):
l = 3.5
pi = 3.14
ratio = 0.47
quad = madseq.Element(None, 'QUADRUPOLE', odicti(K1=pi, L=l))
scaled = madseq.rescale_makethin(quad, ratio)
self.assertEqual(scaled['KNL'], [0, pi*l*ratio])
self.assertEqual(scaled['lrad'], l * ratio)
self.assertEqual(scaled.get('K1'), None)
self.assertEqual(scaled.type, 'multipole')
def test_sbend(self):
l = 3.5
pi = 3.14
ratio = 0.3
sbend = madseq.Element(None, 'SBEND', odicti(angle=pi, hgap=1, L=l))
scaled = madseq.rescale_makethin(sbend, ratio)
self.assertEqual(scaled['KNL'], [pi*ratio])
self.assertEqual(scaled['lrad'], l * ratio)
self.assertEqual(scaled.get('angle'), None)
self.assertEqual(scaled.get('hgap'), None)
self.assertEqual(scaled.type, 'multipole')
def test_eq(self):
el0 = madseq.Element('a', 'b', odicti(c=1))
self.assertEqual(el0, madseq.Element('a', 'b', odicti(c=1)))
self.assertNotEqual(el0, madseq.Element('x', 'b', odicti(c=1)))
self.assertNotEqual(el0, madseq.Element('a', 'x', odicti(c=1)))
self.assertNotEqual(el0, madseq.Element('a', 'b', odicti(x=1)))
self.assertNotEqual(el0, madseq.Element('a', 'b', odicti(c=2)))
def test_detect(self):
input_elements = [
madseq.Element(None, 'outer', None),
madseq.Element('seq', 'sequence', odicti(l=1)),
madseq.Element(None, 'inner', odicti(k=2)),
madseq.Element(None, 'endsequence', odicti()),
madseq.Element(None, 'outer', None)
]
s = list(madseq.Sequence.detect(input_elements))
seq = s[1]
self.assertTrue(s[0] is input_elements[0])
self.assertTrue(seq.head is input_elements[1])
self.assertEqual(len(seq.body), 1)
self.assertTrue(seq.body[0] is input_elements[2])
self.assertTrue(seq.tail is input_elements[3])
self.assertTrue(s[2] is input_elements[4])
self.assertEqual(seq.name, 'seq')
self.assertEqual(str(seq), ("seq: sequence, l=1;\n"
"inner, k=2;\n"
"endsequence;"))
def test_arbitrary(self):
l = 1.3
pi = 3.14
ratio = 0.2
el = madseq.Element('foo', 'arbitrary', odicti(angle=pi, L=l))
scaled = madseq.rescale_thick(el, ratio)
# check that the original element was not modified:
self.assertFalse(scaled is el)
self.assertEqual(el['angle'], pi)
self.assertEqual(el['l'], l)
# check that the scaled element is scaled correctly:
self.assertEqual(scaled['angle'], pi) # don't scale unknown fields!
self.assertEqual(scaled['l'], l * ratio)
self.assertEqual(scaled.type, 'arbitrary')
self.assertEqual(scaled.name, 'foo')
def parse_args(text):
"""Parse argument list into ordered dictionary."""
# TODO: use .split(',') to make expression simpler
return odicti((key, Value.parse(val, assign))
for key,assign,val in regex.arg.findall(text or ''))
def test_format(self):
el = madseq.Element('foo', 'bar', odicti([('a', 1), ('b', 2)]))
self.assertEqual(str(el),
'foo: bar, a=1, b=2;')
def test_setitem(self):
el0 = madseq.Element(None, None, odicti(a=1))
el1 = madseq.Element(None, None, odicti(), el0)
el1['a'] = 2
self.assertEqual(el1['a'], 2)
self.assertEqual(el0['a'], 1)
def test_contains(self):
el = madseq.Element(None, None, odicti(a=1))
self.assertTrue('a' in el)
self.assertFalse('b' in el)
def _getstate(self):
"""Get a serializeable state for :class:`Json` and :class:`Yaml`."""
return odicti([('name', self.name),
('type', self.type)] + list(self.args.items()))