def test_lattice_creation_from_lattice_inherits_attributes():
d = elements.Dipole('d1', 1, BendingAngle=numpy.pi, Energy=5.e+6,
PassMethod='BndMPoleSymplectic4RadPass')
lat1 = Lattice([d], name='lattice', energy=3.e+6, periodicity=32,
an_attr=12)
lat2 = Lattice(lat1, another_attr=5)
assert id(lat1) != id(lat2)
assert len(lat2) == 1
assert lat2.name == 'lattice'
assert lat2.energy == 3.e+6
assert lat2.periodicity == 32
assert lat2.radiation is True
assert lat2.another_attr == 5
with pytest.raises(AttributeError):
assert lat2.an_attr == 12
def test_lattice_energy_radiation_periodicity():
d = elements.Dipole('d1', 1, BendingAngle=numpy.pi/16, Energy=5.e+6,
PassMethod='BndMPoleSymplectic4RadPass')
lat = Lattice([d], name='lattice', energy=3.e+6)
assert lat.energy == 3.e+6
assert lat.periodicity == 32
assert lat.radiation is True
def test_exact_hamiltonian_pass(rin):
drift = elements.Multipole('m1', 1, [0, 0, 0, 0], [0, 0, 0, 0])
drift.Type = 0
drift.PassMethod = 'ExactHamiltonianPass'
drift.BendingAngle = 0
l = Lattice([drift], name='lat', energy=3e9)
atpass(l, rin, 1)
def _fring(ring, split_inds=[], detuning_elem=None):
all_rings, merged_ring = _rearrange(ring, split_inds=split_inds)
ibegs = get_cells(merged_ring, checkname('xbeg'))
iends = get_cells(merged_ring, checkname('xend'))
_, orbit = merged_ring.find_orbit(refpts=ibegs | iends)
if detuning_elem is None:
detuning_elem = gen_detuning_elem(merged_ring, orbit[-1])
else:
detuning_elem.T1 = -orbit[-1]
detuning_elem.T2 = orbit[-1]
fastring = []
for counter, r in enumerate(all_rings):
cavs = [e for e in r if e.PassMethod == 'CavityPass']
[r.remove(c) for c in cavs]
lin_elem = gen_m66_elem(r, orbit[2*counter],
orbit[2*counter+1])
lin_elem.FamName = lin_elem.FamName + '_' + str(counter)
[fastring.append(c) for c in cavs]
fastring.append(lin_elem)
fastring.append(detuning_elem)
try:
qd_elem = gen_quantdiff_elem(merged_ring)
fastring.append(qd_elem)
except ValueError: # No synchrotron radiation => no diffusion element
pass
fastring = Lattice(fastring, **vars(ring))
return fastring
def load_mat(filename, **kwargs):
"""Create a lattice object from a matmab mat-file
PARAMETERS
filename name of a '.mat' file
KEYWORDS
mat_key name of the Matlab variable containing the lattice.
Default: Matlab variable name if there is only one,
otherwise 'RING'
check=True if False, skip the coherence tests
quiet=False If True, suppress the warning for non-standard classes
keep_all=False if True, keep RingParam elements as Markers
name Name of the lattice
(default: taken from the lattice, or '')
energy Energy of the lattice
(default: taken from the elements)
periodicity Number of periods
(default: taken from the elements, or 1)
* all other keywords will be set as Lattice attributes
OUTPUT
Lattice object
"""
if 'key' in kwargs: # process the deprecated 'key' keyword
kwargs.setdefault('mat_key', kwargs.pop('key'))
return Lattice(ringparam_filter,
matfile_generator,
abspath(filename),
iterator=params_filter,
**kwargs)
def test_lattice_creation_gets_attributes_from_arguments():
lat = Lattice(name='lattice', energy=3.e+6, periodicity=32, an_attr=12)
assert len(lat) == 0
assert lat.name == 'lattice'
assert lat.energy == 3.e+6
assert lat.periodicity == 32
assert lat._radiation is False
assert lat.an_attr == 12
def _matlab_scanner(element_list, **kwargs):
"""This function simulates a mat-file reading but replaces the mat-file by
a list of elements"""
latt = Lattice(ringparam_filter,
no_filter,
element_list,
iterator=params_filter,
**kwargs)
return vars(latt)
def test_lattice_voltage_harmonic_number():
rf = elements.RFCavity('rf', 0, 0.2e6, 0.5e9, 5, 3.e6)
d = elements.Dipole('d1', 2.99792458, BendingAngle=numpy.pi/5)
lat = Lattice([rf, d], name='lattice')
assert lat.energy == 3.e+6
assert lat.periodicity == 10
assert lat.rf_voltage == 2e6
assert lat.revolution_frequency == 10.e6
assert lat.harmonic_number == 50
assert lat.radiation is True
def test_lattice_creation_short_scan_reads_radiation_status_correctly():
d = elements.Dipole('d1',
1,
BendingAngle=numpy.pi,
Energy=5.e+6,
PassMethod='BndMPoleSymplectic4RadPass')
lat = Lattice([d], name='lattice', energy=3.e+6, periodicity=32)
assert len(lat) == 1
assert lat.name == 'lattice'
assert lat.energy == 3.e+6
assert lat.periodicity == 32
assert lat._radiation is True
def load_var(matlat, **kwargs):
"""Create a lattice from a Matlab cell array"""
# noinspection PyUnusedLocal
def var_generator(params, latt):
for elem in latt:
yield element_from_dict(elem)
return Lattice(ringparam_filter,
var_generator,
matlat,
iterator=params_filter,
**kwargs)
def test_exact_hamiltonian_pass_with_dls_dipole(rin):
bend = elements.Multipole('rb', 0.15, [0, 0, 0, 0],
[-0.0116333, 3.786786, 0, 0])
bend.Type = 1
bend.PassMethod = 'ExactHamiltonianPass'
bend.BendingAngle = -0.001745
bend.Energy = 3.5e9
bend.MaxOrder = 3
l = Lattice([bend], name='lat', energy=3.5e9)
atpass(l, rin, 1)
# Results from Matlab
expected = numpy.array([9.23965e-9, 1.22319e-5, 0, 0, 0,
-4.8100e-10]).reshape(6, 1)
numpy.testing.assert_allclose(rin, expected, rtol=1e-5, atol=1e-6)
def test_lattice_string_ordering():
lat = Lattice([elements.Drift('D0', 1.0, attr1=numpy.array(0))],
name='lat', energy=5, periodicity=1, attr2=3)
latstr = str(lat)
assert latstr.startswith("Lattice(<1 elements>, name='lat', "
"energy=5, particle=Particle('electron'), "
"periodicity=1")
assert latstr.endswith("attr2=3)")
latrepr = repr(lat)
assert latrepr.startswith("Lattice([Drift('D0', 1.0, attr1=array(0))], "
"name='lat', "
"energy=5, particle=Particle('electron'), "
"periodicity=1")
assert latrepr.endswith("attr2=3)")
def test_lattice_string_ordering():
lat = Lattice([elements.Drift('D0', 1.0, attr1=numpy.array(0))],
name='lat',
energy=5,
periodicity=1,
attr2=3)
# Default dictionary ordering is only in Python >= 3.6
if sys.version_info < (3, 6):
assert lat.__str__().startswith("Lattice(<1 elements>, ")
assert lat.__str__().endswith(", attr2=3)")
assert lat.__repr__().startswith("Lattice([Drift('D0', 1.0, "
"attr1=array(0))], ")
assert lat.__repr__().endswith(", attr2=3)")
else:
assert lat.__str__() == ("Lattice(<1 elements>, name='lat', energy=5,"
" periodicity=1, attr2=3)")
assert lat.__repr__() == ("Lattice([Drift('D0', 1.0, attr1=array(0))],"
" name='lat', energy=5, periodicity=1,"
" attr2=3)")
def load_tracy(filename, **kwargs):
try:
harmonic_number = kwargs.pop("harmonic_number")
lattice_key = kwargs.pop("lattice_key", "cell")
def elem_iterator(params, tracy_file):
with open(params.setdefault("tracy_file", tracy_file)) as f:
contents = f.read()
element_lines, energy = expand_tracy(contents, lattice_key,
harmonic_number)
params.setdefault("energy", energy)
for line in element_lines:
yield line
return Lattice(abspath(filename), iterator=elem_iterator, **kwargs)
except Exception as e:
raise ValueError("Failed to load tracy lattice {}: {}".format(
filename, e))
def _rearrange(ring, split_inds=[]):
inds = numpy.append(split_inds, [0, len(ring)+1])
inds = numpy.unique(inds)
all_rings = [ring[int(b):int(e)] for b, e in zip(inds[:-1], inds[1:])]
ringm = []
for ring_slice in all_rings:
ring_slice.insert(0, Marker('xbeg'))
ring_slice.append(Marker('xend'))
cavs = [e for e in ring_slice if isinstance(e, RFCavity)]
newpass = ['IdentityPass' if c.Length == 0
else 'DriftPass' for c in cavs]
for c, pm in zip(cavs, newpass):
c.PassMethod = pm
uni_freq = numpy.unique([e.Frequency for e in cavs])
for fr in numpy.atleast_1d(uni_freq):
cavf = [c for c in cavs if c.Frequency == fr]
vol = numpy.sum([c.Voltage for c in cavf])
cavl = RFCavity('CAVL', 0, vol, fr,
cavf[0].HarmNumber, cavf[0].Energy)
cavl.TimeLag = cavf[0].TimeLag
ring_slice.append(cavl)
ringm = ringm + ring_slice
return all_rings, Lattice(ringm, energy=ring.energy)
def load_m(filename, **kwargs):
"""Create a lattice object from a matlab m-file
PARAMETERS
filename name of a '.m' file
KEYWORDS
keep_all=False if True, keep RingParam elements as Markers
name Name of the lattice
(default: taken from the elements, or '')
energy Energy of the lattice
(default: taken from the elements)
periodicity Number of periods
(default: taken from the elements, or 1)
* all other keywords will be set as Lattice attributes
OUTPUT
Lattice object
"""
return Lattice(ringparam_filter,
mfile_generator,
abspath(filename),
iterator=params_filter,
**kwargs)
def load_repr(filename, **kwargs):
"""Load a Lattice object stored as a repr-file
PARAMETERS
filename name of the '.repr' file
KEYWORDS
name Name of the lattice (default: taken from the file)
energy Energy of the lattice (default: taken from the file)
periodicity Number of periods (default: taken from the file)
* all other keywords will be set as Lattice attributes
OUTPUT
Lattice object
"""
def elem_iterator(params, repr_file):
with open(params.setdefault('repr_file', repr_file), 'rt') as file:
# the 1st line is the dictionary of saved lattice parameters
for k, v in eval(next(file)).items():
params.setdefault(k, v)
for line in file:
yield element_from_string(line.strip())
return Lattice(abspath(filename), iterator=elem_iterator, **kwargs)
def test_gwig_symplectic_pass(rin, passmethod):
# Parameters copied from one of the Diamond wigglers.
wiggler = elements.Wiggler('w', 1.15, 0.05, 0.8, 3e9)
wiggler.PassMethod = passmethod
l = Lattice([wiggler], name='lat', energy=3e9)
atpass(l, rin, 1)
def test_bndstrmpole_symplectic_4_pass(rin):
bend = elements.Dipole('b', 1.0)
bend.PassMethod = 'BndStrMPoleSymplectic4Pass'
l = Lattice([bend], name='lat', energy=3e9)
atpass(l, rin, 1)
def test_item_is_not_an_AT_element_warns_correctly():
with pytest.warns(AtWarning):
Lattice(['a'], energy=0, periodicity=1)
def test_lattice_energy_is_not_defined_raises_AtError():
with pytest.raises(AtError):
Lattice()
def test_lattice_energy_is_not_defined_raises_AtError():
d = elements.Dipole('d1', 1, BendingAngle=numpy.pi)
with pytest.raises(AtError):
Lattice([d])