def test_lattice_without_symmetry():
lat = Lattice("")
assert lat.cell_length is None
assert lat.cell_bounds is None
lat = Lattice("", 6)
assert lat.cell_length is None
assert lat.cell_bounds is None
def test_get_and_set_unitconv():
lat = Lattice("")
with pytest.raises(KeyError):
lat._data_source_manager._uc["field1"]
uc = mock.Mock()
lat.set_unitconv("field1", uc)
assert lat._data_source_manager._uc["field1"] == uc
assert lat.get_unitconv("field1") == uc
def test_get_element_devices_raises_ValueError_for_mismatched_family(
simple_lattice):
with pytest.raises(ValueError):
devices = simple_lattice.get_element_devices("not-a-family", "x")
basic_element = simple_lattice.get_elements("family")[0]
basic_lattice = Lattice("basic_lattice")
del basic_element._data_source_manager._data_sources[pytac.LIVE]
basic_lattice.add_element(basic_element)
devices = basic_lattice.get_element_devices("family", "x")
assert devices == []
def test_set_lattice_reference():
elem1 = Element(1.0, "BPM")
lat1 = Lattice("one")
elem2 = Element(2.0, "BPM", lattice=lat1)
lat2 = Lattice("two")
assert elem1._lattice is None
assert elem2._lattice == lat1
elem1.set_lattice(lat1)
elem2.set_lattice(lat2)
assert elem1._lattice == lat1
assert elem2._lattice == lat2
def load_unitconv(mode_dir: Path, lattice: Lattice) -> None:
"""Load the unit conversion objects from a file.
Args:
mode_dir: Path to directory containing CSV files.
lattice: The lattice object that will be used.
"""
unitconvs: Dict[int, UnitConv] = {}
# Assemble datasets from the polynomial file
unitconvs.update(load_poly_unitconv(mode_dir / POLY_FILENAME))
# Assemble datasets from the pchip file
unitconvs.update(load_pchip_unitconv(mode_dir / PCHIP_FILENAME))
# Add the unitconv objects to the elements
with csv_loader(mode_dir / UNITCONV_FILENAME) as csv_reader:
for item in csv_reader:
# Special case for element 0: the lattice itself.
if int(item["el_id"]) == 0:
if item["uc_type"] != "null":
# Each element needs its own unitconv object as
# it may for example have different limit.
uc = copy.copy(unitconvs[int(item["uc_id"])])
uc.phys_units = item["phys_units"]
uc.eng_units = item["eng_units"]
upper, lower = (
float(lim) if lim != "" else None
for lim in [item["upper_lim"], item["lower_lim"]])
uc.set_conversion_limits(lower, upper)
else:
uc = NullUnitConv(item["eng_units"], item["phys_units"])
lattice.set_unitconv(item["field"], uc)
else:
element = lattice[int(item["el_id"]) - 1]
# For certain magnet types, we need an additional rigidity
# conversion factor as well as the raw conversion.
if item["uc_type"] == "null":
uc = NullUnitConv(item["eng_units"], item["phys_units"])
else:
# Each element needs its own unitconv object as
# it may for example have different limit.
uc = copy.copy(unitconvs[int(item["uc_id"])])
if any(
element.is_in_family(f)
for f in ("HSTR", "VSTR", "Quadrupole",
"Sextupole", "Bend")):
energy = lattice.get_value("energy", units=pytac.PHYS)
uc.set_post_eng_to_phys(utils.get_div_rigidity(energy))
uc.set_pre_phys_to_eng(utils.get_mult_rigidity(energy))
uc.phys_units = item["phys_units"]
uc.eng_units = item["eng_units"]
upper, lower = (
float(lim) if lim != "" else None
for lim in [item["upper_lim"], item["lower_lim"]])
uc.set_conversion_limits(lower, upper)
element.set_unitconv(item["field"], uc)
def test_element_properties_with_lattice():
e1 = Element(3.1, "DRFIT", "d1")
e2 = Element(1.3, "DRFIT", "d2")
lat = Lattice("", symmetry=2)
lat.add_element(e1)
lat.add_element(e2)
assert e1.index == 1
assert e2.index == 2
assert e1.s == 0.0
assert e2.s == 3.1
assert e1.cell == 1
assert e2.cell == 2
def test_element_representation():
elem = Element(0.1, "BPM")
assert str(elem) == "<Element length 0.1 m, families >"
elem.add_to_family("fam1")
assert str(elem) == "<Element length 0.1 m, families fam1>"
elem.name = "bpm1"
assert str(elem) == "<Element 'bpm1', length 0.1 m, families fam1>"
lat = Lattice("")
lat.add_element(elem)
assert str(elem) == ("<Element 'bpm1', index 1, length 0.1 m, families " "fam1>")
lat.symmetry = 2
assert str(elem) == (
"<Element 'bpm1', index 1, length 0.1 m, cell 1, " "families fam1>"
)
elem.name = None
assert str(elem) == ("<Element index 1, length 0.1 m, cell 1, families " "fam1>")
def test_create_element():
lat = Lattice("")
e = Element(6.0, "BPM", "bpm1", lat)
assert e.length == 6.0
assert e.type_ == "BPM"
assert e.name == "bpm1"
assert e._lattice == lat
def test_add_element_to_lattice():
lat1 = Lattice("lat1")
elem = Element(0.5, "DRIFT")
assert lat1._elements == []
assert elem._lattice is None
lat1.add_element(elem)
assert lat1._elements == [elem]
assert elem._lattice == lat1
lat2 = Lattice("lat2")
lat2.add_element(elem)
assert elem._lattice == lat2
def test_lattice_cell_properties():
lat = Lattice("", 2)
for i in range(5):
lat.add_element(Element(0.5, "DRIFT"))
assert lat.cell_length == 1.25
assert lat.cell_bounds == [1, 4, 5]
def test_create_lattice():
lat = Lattice(LATTICE_NAME)
assert (len(lat)) == 0
assert lat.name == LATTICE_NAME
def simple_lattice(simple_element, x_device, y_device, mock_sim_data_source,
unit_uc, double_uc):
lattice = Lattice(LATTICE_NAME, symmetry=6)
lattice.add_element(simple_element)
lattice.set_data_source(DeviceDataSource(), pytac.LIVE)
lattice.add_device("x", x_device, unit_uc)
lattice.add_device("y", y_device, double_uc)
lattice.set_data_source(mock_sim_data_source, pytac.SIM)
return lattice