Merge pull request #54 from DiamondLightSource/diamond2-conversion

Diamond2 conversion:
-Make atip compatible with Diamond 2 lattice.
-Support multipoles in atip
-Make create_lattice_matfile.m compatible with Diamond 2 lattice and improve code readibility
-Regenerate all lattice files using this new script.
-Add Diamond 2 lattice to tests
-Made a few improvements to the tests

Author: ptsOSL

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(

src/atip/rings/48.mat

@@ -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

src/atip/rings/DIAD.mat

@@ -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,21 +202,21 @@ 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.
 
         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)

src/atip/rings/I04.mat

@@ -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()

src/atip/rings/create_lattice_matfile.m

@@ -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()