diff --git a/src/atip/load_sim.py b/src/atip/load_sim.py index 9aa70d3..698f66e 100644 --- a/src/atip/load_sim.py +++ b/src/atip/load_sim.py @@ -9,7 +9,7 @@ from atip.simulator import ATSimulator # List of all the element fields that can be currently simulated. -SIMULATED_FIELDS = {"a1", "b0", "b1", "b2", "x", "y", "f", "x_kick", "y_kick"} +SIMULATED_FIELDS = {"a1", "b0", "b1", "b2", "b3", "x", "y", "f", "x_kick", "y_kick"} def load_from_filepath( diff --git a/src/atip/rings/48.mat b/src/atip/rings/48.mat new file mode 100644 index 0000000..4436653 Binary files /dev/null and b/src/atip/rings/48.mat differ diff --git a/src/atip/rings/DIAD.mat b/src/atip/rings/DIAD.mat index a424136..3d73778 100644 Binary files a/src/atip/rings/DIAD.mat and b/src/atip/rings/DIAD.mat differ diff --git a/src/atip/rings/I04.mat b/src/atip/rings/I04.mat index df9b1cc..895f914 100644 Binary files a/src/atip/rings/I04.mat and b/src/atip/rings/I04.mat differ diff --git a/src/atip/rings/create_lattice_matfile.m b/src/atip/rings/create_lattice_matfile.m index 47cd4a7..4021854 100644 --- a/src/atip/rings/create_lattice_matfile.m +++ b/src/atip/rings/create_lattice_matfile.m @@ -1,70 +1,75 @@ function create_lattice_matfile(filename) % Creates a .mat file AT lattice compatible with ATIP. % If a filename is given that file will be updated to ATIP standard. Otherwise -% the ring is taken from either of the 'RING' or 'THERING' global variables, -% with 'RING' taking priority. If a filename is not passed the updated lattice -% will be stored in 'lattice.mat'. +% ATIP_RING is initially taken from 'THERING' global variable and save as +% 'lattice.mat'. if ~(nargin == 0) - load(filename, 'RING'); + load(filename, 'ATIP_RING'); end - if ~exist('RING', 'var') - global RING; - if isempty(RING) - global THERING; - RING = THERING; + + if ~exist('ATIP_RING', 'var') + global THERING; + if isempty(THERING) + disp('THERING global variable is empty, try running storageringinit(Ringmode). Exiting with error.'); + exit(1) + else + ATIP_RING=THERING + disp('Using global THERING and saving it to global ATIP_RING.'); end + else + disp('Using loaded ATIP_RING from file.'); end - if isempty(RING) - disp('Unable to load a ring from file or global variables.'); - return; - end + + fprintf('Initial lattice has dimensions: %s\n', mat2str(size(ATIP_RING))) % Correct dimension order if necessary. - if size(RING, 1) == 1 - RING = permute(RING, [2 1]); + if size(ATIP_RING, 1) == 1 + ATIP_RING = permute(ATIP_RING, [2 1]); end + % Correct classes and pass methods. - for x = 1:length(RING) - if strcmp(RING{x, 1}.FamName, 'BPM10') + for x = 1:length(ATIP_RING) + if strcmp(ATIP_RING{x, 1}.FamName, 'BPM10') % Wouldn't be correctly classed by class guessing otherwise. - RING{x, 1}.Class = 'Monitor'; - elseif (strcmp(RING{x, 1}.FamName, 'HSTR') || strcmp(RING{x, 1}.FamName, 'VSTR')) - RING{x, 1}.Class = 'Corrector'; - elseif (strcmp(RING{x, 1}.FamName, 'HTRIM') || strcmp(RING{x, 1}.FamName, 'VTRIM')) - RING{x, 1}.Class = 'Corrector'; + ATIP_RING{x, 1}.Class = 'Monitor'; + elseif (strcmp(ATIP_RING{x, 1}.FamName, 'HSTR') || strcmp(ATIP_RING{x, 1}.FamName, 'VSTR')) + ATIP_RING{x, 1}.Class = 'Corrector'; + elseif (strcmp(ATIP_RING{x, 1}.FamName, 'HTRIM') || strcmp(ATIP_RING{x, 1}.FamName, 'VTRIM')) + ATIP_RING{x, 1}.Class = 'Corrector'; end - if isfield(RING{x, 1}, 'Class') - if strcmp(RING{x, 1}.Class, 'SEXT') - RING{x, 1}.Class = 'Sextupole'; + + if isfield(ATIP_RING{x, 1}, 'Class') + if strcmp(ATIP_RING{x, 1}.Class, 'SEXT') + ATIP_RING{x, 1}.Class = 'Sextupole'; end end - if strcmp(RING{x, 1}.PassMethod, 'ThinCorrectorPass') - % ThinCorrectorPass no longer exists in AT. - RING{x, 1}.PassMethod = 'CorrectorPass'; - elseif strcmp(RING{x, 1}.PassMethod, 'GWigSymplecticPass') - RING{x, 1}.Class = 'Wiggler'; + + if strcmp(ATIP_RING{x, 1}.PassMethod, 'GWigSymplecticPass') + ATIP_RING{x, 1}.Class = 'Wiggler'; end end - % Remove elements. Done this way because the size of RING changes during - % the loop. + % Remove elements. Done this way because the size of ATIP_RING changes + % during the loop. y = 1; - while y < length(RING) - % I should probably transfer the attributes of the deleted corrector - % elements to the sextupole but cba. - if (strcmp(RING{y, 1}.FamName, 'HSTR') && strcmp(RING{y-1, 1}.Class, 'Sextupole')) - RING(y, :) = []; % Delete hstrs that are preceded by a sextupole. - elseif (strcmp(RING{y, 1}.FamName, 'VSTR') && strcmp(RING{y-1, 1}.Class, 'Sextupole')) - RING(y, :) = []; % Delete vstrs that are preceded by a sextupole. + while y < length(ATIP_RING) + % The data within the deleted elements is not needed + if strcmp(ATIP_RING{y, 1}.FamName, 'HSTR') && ATIP_RING{y, 1}.Length == 0 && (strcmp(ATIP_RING{y-1, 1}.Class, 'Sextupole') || strcmp(ATIP_RING{y-1, 1}.Class, 'Multipole')) + ATIP_RING(y, :) = []; % Delete hstrs that are preceded by a sextupole or multipole. + elseif strcmp(ATIP_RING{y, 1}.FamName, 'VSTR') && ATIP_RING{y, 1}.Length == 0 && (strcmp(ATIP_RING{y-1, 1}.Class, 'Sextupole') || strcmp(ATIP_RING{y-1, 1}.Class, 'Multipole')) + ATIP_RING(y, :) = []; % Delete vstrs that are preceded by a sextupole or multipole. else y = y + 1; end end - if isfield(RING{1, 1}, 'TwissData') - RING{1, 1} = rmfield(RING{1, 1}, 'TwissData'); + + if isfield(ATIP_RING{1, 1}, 'TwissData') + ATIP_RING{1, 1} = rmfield(ATIP_RING{1, 1}, 'TwissData'); end + + fprintf('Converted ATIP_RING has dimensions: %s\n', mat2str(size(ATIP_RING))) if nargin == 0 - save('lattice.mat', 'RING'); + save('lattice.mat', 'ATIP_RING'); else - save(filename, 'RING'); + save(filename, 'ATIP_RING'); end end diff --git a/src/atip/sim_data_sources.py b/src/atip/sim_data_sources.py index cb88528..b1e7226 100644 --- a/src/atip/sim_data_sources.py +++ b/src/atip/sim_data_sources.py @@ -72,6 +72,7 @@ def __init__(self, at_element, index, atsim, fields=None): "a1": partial(self._get_PolynomA, 1), "b1": partial(self._get_PolynomB, 1), "b2": partial(self._get_PolynomB, 2), + "b3": partial(self._get_PolynomB, 3), "b0": self._get_BendingAngle, "f": self._get_Frequency, } @@ -81,6 +82,7 @@ def __init__(self, at_element, index, atsim, fields=None): "a1": partial(self._set_PolynomA, 1), "b1": partial(self._set_PolynomB, 1), "b2": partial(self._set_PolynomB, 2), + "b3": partial(self._set_PolynomB, 3), "b0": self._set_BendingAngle, "f": self._set_Frequency, } @@ -200,13 +202,13 @@ def _get_KickAngle(self, cell): """A data handling function used to get the value of a specific cell of the KickAngle attribute of the AT element. - .. Note:: If the Corrector is attached to a Sextupole then KickAngle - needs to be returned from cell 0 of the applicable Polynom(A/B) - attribute and so a conversion must take place. For independent - Correctors KickAngle can be returned directly from the element's - KickAngle attribute without any conversion. This is because - independent Correctors have a KickAngle attribute in our AT lattice, - but those attached to Sextupoles do not. + .. Note:: If the Corrector is attached to a Sextupole or Octupole then + KickAngle needs to be returned from cell 0 of the applicable Polynom(A/B) + attribute and so a conversion must take place. For independent + Correctors KickAngle can be returned directly from the element's + KickAngle attribute without any conversion. This is because + independent Correctors have a KickAngle attribute in our AT lattice, + but those attached to Sextupoles do not. Args: cell (int): Which cell of KickAngle to get. @@ -214,7 +216,7 @@ def _get_KickAngle(self, cell): Returns: float: The kick angle of the specified cell. """ - if isinstance(self._at_element, at.elements.Sextupole): + if isinstance(self._at_element, (at.elements.Sextupole, at.elements.Multipole)): length = self._at_element.Length if cell == 0: return -(self._at_element.PolynomB[0] * length) @@ -227,19 +229,19 @@ def _set_KickAngle(self, cell, value): """A data handling function used to set the value of a specific cell of the KickAngle attribute of the AT element. - .. Note:: If the Corrector is attached to a Sextupole then KickAngle - needs to be assigned to cell 0 of the applicable Polynom(A/B) - attribute and so a conversion must take place. For independent - Correctors KickAngle can be assigned directly to the element's - KickAngle attribute without any conversion. This is because - independent Correctors have a KickAngle attribute in our AT lattice, - but those attached to Sextupoles do not. + .. Note:: If the Corrector is attached to a Sextupole or Octupole then + KickAngle needs to be assigned to cell 0 of the applicable Polynom(A/B) + attribute and so a conversion must take place. For independent + Correctors KickAngle can be assigned directly to the element's + KickAngle attribute without any conversion. This is because + independent Correctors have a KickAngle attribute in our AT lattice, + but those attached to Sextupoles do not. Args: cell (int): Which cell of KickAngle to set. value (float): The angle to be set. """ - if isinstance(self._at_element, at.elements.Sextupole): + if isinstance(self._at_element, (at.elements.Sextupole, at.elements.Multipole)): length = self._at_element.Length if cell == 0: self._at_element.PolynomB[0] = -(value / length) diff --git a/tests/conftest.py b/tests/conftest.py index b898298..4add592 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,25 +62,31 @@ def atlds(): return atip.sim_data_sources.ATLatticeDataSource(mock.Mock()) -@pytest.fixture() -def at_lattice(): - return atip.utils.load_at_lattice("I04") +@pytest.fixture(scope="function", params=["I04"]) +def at_and_pytac_lattices(request): + lattices = [] + lattices.append(load_csv.load(request.param, cs.ControlSystem())) + lattices.append(atip.utils.load_at_lattice(request.param)) + return lattices -@pytest.fixture(scope="session") -def pytac_lattice(): - return load_csv.load("DIAD", cs.ControlSystem()) +@pytest.fixture(scope="function", params=["I04"]) +def pytac_lattice(request): + return load_csv.load(request.param, cs.ControlSystem()) -@pytest.fixture(scope="session") -def mat_filepath(): - here = os.path.dirname(__file__) - return os.path.realpath(os.path.join(here, "../src/atip/rings/DIAD.mat")) +@pytest.fixture(scope="function", params=["I04"]) +def at_lattice(request): + return atip.utils.load_at_lattice(request.param) -@pytest.fixture(scope="session") -def at_diad_lattice(mat_filepath): - return at.load.load_mat(mat_filepath) +@pytest.fixture(scope="function", params=["DIAD"]) +def lattice_filepath(request): + here = os.path.dirname(__file__) + filepath = os.path.realpath( + os.path.join(here, f"../src/atip/rings/{request.param}.mat") + ) + return filepath @pytest.fixture() diff --git a/tests/test_load.py b/tests/test_load.py index 91a5542..3929bc1 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -5,9 +5,21 @@ import atip +RINGMODES_TO_TEST = ["I04", "DIAD", "48"] -def test_load_pytac_side(pytac_lattice, at_diad_lattice): - lat = atip.load_sim.load(pytac_lattice, at_diad_lattice) + +@pytest.mark.parametrize( + "at_lattice", + RINGMODES_TO_TEST, + indirect=True, +) +def test_load_atip_lattice(request, at_lattice): + assert at_lattice.name == request.node.callspec.params["at_lattice"] + + +@pytest.mark.parametrize("at_and_pytac_lattices", RINGMODES_TO_TEST, indirect=True) +def test_load_pytac_side(at_and_pytac_lattices): + lat = atip.load_sim.load(at_and_pytac_lattices[0], at_and_pytac_lattices[1]) # Check lattice has simulator data source assert pytac.SIM in lat._data_source_manager._data_sources # Check all elements have simulator data source @@ -18,22 +30,29 @@ def test_load_pytac_side(pytac_lattice, at_diad_lattice): assert isinstance(lat._data_source_manager._uc["mu"], pytac.units.NullUnitConv) -def test_load_from_filepath(pytac_lattice, mat_filepath): - atip.load_sim.load_from_filepath(pytac_lattice, mat_filepath) +@pytest.mark.parametrize( + ["pytac_lattice", "lattice_filepath"], + [(mode, mode) for mode in RINGMODES_TO_TEST], + indirect=True, +) +def test_load_atip_and_pytac_lattices(pytac_lattice, lattice_filepath): + atip.load_sim.load_from_filepath(pytac_lattice, lattice_filepath) -def test_load_with_non_callable_callback_raises_TypeError( - pytac_lattice, at_diad_lattice -): +@pytest.mark.parametrize("at_and_pytac_lattices", RINGMODES_TO_TEST, indirect=True) +def test_load_with_non_callable_callback_raises_TypeError(at_and_pytac_lattices): with pytest.raises(TypeError): - atip.load_sim.load(pytac_lattice, at_diad_lattice, "") + atip.load_sim.load(at_and_pytac_lattices[0], at_and_pytac_lattices[1], "") -def test_load_with_callback(pytac_lattice, at_diad_lattice): +@pytest.mark.parametrize("at_and_pytac_lattices", RINGMODES_TO_TEST, indirect=True) +def test_load_with_callback(at_and_pytac_lattices): callback_func = mock.Mock() - lat = atip.load_sim.load(pytac_lattice, at_diad_lattice, callback_func) + lat = atip.load_sim.load( + at_and_pytac_lattices[0], at_and_pytac_lattices[1], callback_func + ) atsim = lat._data_source_manager._data_sources[pytac.SIM]._atsim - atip.utils.trigger_calc(pytac_lattice) + atip.utils.trigger_calc(at_and_pytac_lattices[0]) atsim.wait_for_calculations() callback_func.assert_called_once_with()