"""Module responsible for handling the loading of simulator data sources."""
import at
import pytac
from pytac.exceptions import FieldException
from pytac.load_csv import DEFAULT_UC
from atip.simulator import ATSimulator
from atip.sim_data_sources import ATElementDataSource, ATLatticeDataSource
# List of all the element fields that can be currently simulated.
SIMULATED_FIELDS = {"a1", "b0", "b1", "b2", "x", "y", "f", "x_kick", "y_kick"}
[docs]def load_from_filepath(pytac_lattice, at_lattice_filepath, callback=None):
"""Load simulator data sources onto the lattice and its elements.
Args:
pytac_lattice (pytac.lattice.Lattice): An instance of a Pytac lattice.
at_lattice_filepath (str): The path to a .mat file from which the
Accelerator Toolbox lattice can be loaded.
callback (callable): To be called after completion of each round of
physics calculations.
Returns:
pytac.lattice.Lattice: The same Pytac lattice object, but now with a
simulator data source fully loaded onto it.
"""
at_lattice = at.load.load_mat(
at_lattice_filepath,
name=pytac_lattice.name,
energy=pytac_lattice.get_value("energy"),
)
return load(pytac_lattice, at_lattice, callback)
[docs]def load(pytac_lattice, at_lattice, callback=None):
"""Load simulator data sources onto the lattice and its elements.
Args:
pytac_lattice (pytac.lattice.Lattice): An instance of a Pytac lattice.
at_lattice (at.lattice_object.Lattice): An instance of an Accelerator
Toolbox lattice object.
callback (callable): To be called after completion of each round of
physics calculations.
Returns:
pytac.lattice.Lattice: The same Pytac lattice object, but now with a
simulator data source fully loaded onto it.
"""
if len(at_lattice) != len(pytac_lattice):
raise ValueError(
"Incompatible AT and Pytac lattices, length mismatch "
"(AT:{0} Pytac:{1}).".format(len(at_lattice), len(pytac_lattice))
)
# Initialise an instance of the ATSimulator Object.
atsim = ATSimulator(at_lattice, callback)
# Set the simulator data source on the Pytac lattice.
pytac_lattice.set_data_source(ATLatticeDataSource(atsim), pytac.SIM)
# Load the sim onto each element.
for e in pytac_lattice:
# Determine which fields each simulated element should have.
sim_fields = list(set(e.get_fields()[pytac.LIVE]) & SIMULATED_FIELDS)
# Set the simulator data source on each element.
e.set_data_source(
ATElementDataSource(at_lattice[e.index - 1], e.index, atsim, sim_fields),
pytac.SIM,
)
# Give any lattice fields not on the live machine a unit conversion object.
for field in pytac_lattice.get_fields()[pytac.SIM]:
try:
pytac_lattice.get_unitconv(field)
except FieldException:
pytac_lattice.set_unitconv(field, DEFAULT_UC)
return pytac_lattice