Transfer map is now method of the MagneticLattice

Author: sergey-tomin

ocelot/cpbd/magnetic_lattice.py

@@ -7,6 +7,7 @@
 from ocelot.cpbd.latticeIO import LatticeIO
 from ocelot.cpbd.transformations.transfer_map import TransferMap
 from ocelot.cpbd.optics import lattice_transfer_map
+from ocelot.cpbd.tm_utils import transfer_maps_mult
 
 import logging
 import re
@@ -103,11 +104,8 @@ def merger(lat, remaining_types=None, remaining_elems=None, init_energy=0.):
         else:
             delta_e = np.sum([tm.get_delta_e() for elem in elem_list for tm in elem.tms])
             lattice = MagneticLattice(elem_list, method=lat.method)
-            R = lattice_transfer_map(lattice, energy=E)
             m = Matrix()
-            m.r = lattice.R
-            m.t = lattice.T
-            m.b = lattice.B
+            m.b, m.r, m.t = lattice.transfer_maps(energy=E)
             m.l = lattice.totalLen
             m.delta_e = delta_e
             E += delta_e
@@ -117,6 +115,7 @@ def merger(lat, remaining_types=None, remaining_elems=None, init_energy=0.):
     _logger.debug("element numbers after: " + str(len(new_lat.sequence)))
     return new_lat
 
+
 def flatten(iterable: Iterator[Any]) -> Generator[Any, None, None]:
     """Flatten arbitrarily nested iterable.
     Special case for strings that avoids infinite recursion.  Non iterables passed
@@ -280,6 +279,23 @@ def save_as_py_file(self, file_name: str, remove_rep_drifts=True, power_supply=F
         LatticeIO.save_lattice(self, tws0=None, file_name=file_name, remove_rep_drifts=remove_rep_drifts,
                                power_supply=power_supply)
 
+    def transfer_maps(self, energy):
+        """
+        Function calculates transfer maps, the first and second orders (R, T), for the whole lattice.
+
+        :param energy: the initial electron beam energy [GeV]
+        :return: B, R, T - matrices
+        """
+        Ra = np.eye(6)
+        Ta = np.zeros((6, 6, 6))
+        Ba = np.zeros((6, 1))
+        E = energy
+        for elem in self.sequence:
+            for Rb, Bb, Tb, tm in zip(elem.R(E), elem.B(E), elem.T(E), elem.tms):
+                Ba, Ra, Ta = transfer_maps_mult(Ba, Ra, Ta, Bb, Rb, Tb)
+                E += tm.get_delta_e()
+        return Ba, Ra, Ta
+
 
 class EndElements:
     suffix_1 = "_1"

ocelot/cpbd/optics.py

@@ -84,19 +84,9 @@ def lattice_transfer_map(lattice, energy):
     :param energy: the initial electron beam energy [GeV]
     :return: R - matrix
     """
-
-    Ra = np.eye(6)
-    Ta = np.zeros((6, 6, 6))
-    Ba = np.zeros((6, 1))
-    E = energy
-    for elem in lattice.sequence:
-        for Rb, Bb, Tb, tm in zip(elem.R(E), elem.B(E), elem.T(E), elem.tms):
-            Ba, Ra, Ta = transfer_maps_mult(Ba, Ra, Ta, Bb, Rb, Tb)
-            #Ba = np.dot(Rb, Ba) + Bb
-            E += tm.get_delta_e()
+    Ba, Ra, Ta = lattice.transfer_maps(energy)
 
     # TODO: Adding Attributes at runtime should be avoided
-    lattice.E = E
     lattice.T_sym = Ta
     lattice.T = Ta #unsym_matrix(deepcopy(Ta))
     lattice.R = Ra

unit_tests/ebeam_test/io_lattice/io_lattice_test.py

@@ -15,7 +15,7 @@
 import copy
 
 
-def test_original_lattice_transfer_map(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_original_lattice_transfer_map(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
 
     r_matrix = lattice_transfer_map(lattice, tws0.E)
@@ -29,7 +29,7 @@ def test_original_lattice_transfer_map(lattice, tws0, method, parametr=None, upd
     assert check_result(result)
 
 
-def test_lattice_save_as_py_file(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_lattice_save_as_py_file(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
 
     lattice.save_as_py_file(file_name="tmp_lattice.py")
@@ -48,7 +48,7 @@ def test_lattice_save_as_py_file(lattice, tws0, method, parametr=None, update_re
     assert check_result(res)
 
 
-def test_lattice_save_as_py_file_w_coupler(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_lattice_save_as_py_file_w_coupler(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     lattice0 = copy.deepcopy(lattice)
     for elem in lattice0.sequence:
@@ -87,7 +87,7 @@ def test_lattice_save_as_py_file_w_coupler(lattice, tws0, method, parametr=None,
 
     assert check_result(res)
 
-def test_original_twiss(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_original_twiss(lattice, tws0, method, parameter=None, update_ref_values=False):
     """Twiss parameters calculation function test"""
 
     tws = twiss(lattice, tws0, nPoints=None)
@@ -103,8 +103,8 @@ def test_original_twiss(lattice, tws0, method, parametr=None, update_ref_values=
     assert check_result(result)
 
 
-@pytest.mark.parametrize('parametr', [False, True])
-def test_lat2input(lattice, tws0, method, parametr, update_ref_values=False):
+@pytest.mark.parametrize('parameter', [False, True])
+def test_lat2input(lattice, tws0, method, parameter, update_ref_values=False):
     """lat2input with tws0 saving function test"""
 
     lines_arr = LatticeIO.lat2input(lattice, tws0=tws0)
@@ -114,13 +114,13 @@ def test_lat2input(lattice, tws0, method, parametr, update_ref_values=False):
     try:
         exec(lines, globals(), loc_dict)
     except Exception as err:
-        assert check_result(['Exception error during the lattice file execution, parametr is ' + str(parametr)])
+        assert check_result(['Exception error during the lattice file execution, parameter is ' + str(parameter)])
 
-    if parametr:
+    if parameter:
         if "tws0" in loc_dict:
             tws0_new = loc_dict['tws0']
         else:
-            assert check_result(['No tws0 in the lattice file, parametr is ' + str(parametr)])
+            assert check_result(['No tws0 in the lattice file, parameter is ' + str(parameter)])
     else:
         tws0_new = tws0
 
@@ -130,7 +130,7 @@ def test_lat2input(lattice, tws0, method, parametr, update_ref_values=False):
         lattice_new_transfer_map_check(lattice_new, tws0_new)
         twiss_new_check(lattice_new, tws0_new)
     else:
-        assert check_result(['No cell variable in the lattice file, parametr is ' + str(parametr)])
+        assert check_result(['No cell variable in the lattice file, parameter is ' + str(parameter)])
 
 
 def lattice_new_transfer_map_check(lattice, tws0):
@@ -142,6 +142,21 @@ def lattice_new_transfer_map_check(lattice, tws0):
     result = check_matrix(r_matrix, r_matrix_ref, TOL, assert_info=' r_matrix for new lattice - ')
     assert check_result(result)
 
+@pytest.mark.parametrize('parameter', [0, 1, 2])
+def test_lattice_transfer_maps_check(lattice, tws0, method, parameter, update_ref_values=False):
+    matrices = lattice.transfer_maps(tws0.E)
+
+
+    #r_matrix_ref = json2numpy(json_read(REF_RES_DIR + 'test_original_lattice_transfer_map.json'))
+    m = matrices[parameter]
+    if update_ref_values:
+        return numpyBRT2json(m)
+
+    m_ref = json2numpyBRT(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + str(parameter) + '.json'))
+
+    result = check_matrix(m, m_ref, TOL, assert_info=' B_matrix for the lattice - ')
+
+    assert check_result(result)
 
 def twiss_new_check(lattice, tws0):
 
@@ -155,7 +170,7 @@ def twiss_new_check(lattice, tws0):
     assert check_result(result)
 
 
-def test_merger(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad")
@@ -185,7 +200,7 @@ def test_merger(lattice, tws0, method, parametr=None, update_ref_values=False):
     assert check_result(result + result2)
 
 
-def test_merger_elem(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger_elem(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad")
@@ -215,7 +230,7 @@ def test_merger_elem(lattice, tws0, method, parametr=None, update_ref_values=Fal
     assert check_result(result + result2)
 
 
-def test_merger_elem_w_coupler(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger_elem_w_coupler(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad")
@@ -252,7 +267,7 @@ def test_merger_elem_w_coupler(lattice, tws0, method, parametr=None, update_ref_
     result2 = check_matrix(lat.T, new_lat.T, TOL, assert_info=' t_matrix - ')
     assert check_result(result + result2)
 
-def test_merger_type(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger_type(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad")
@@ -284,7 +299,7 @@ def test_merger_type(lattice, tws0, method, parametr=None, update_ref_values=Fal
     assert check_result(result + result2)
 
 
-def test_merger_extensive(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger_extensive(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
 
 
@@ -300,7 +315,7 @@ def test_merger_extensive(lattice, tws0, method, parametr=None, update_ref_value
     assert check_result(result + result2)
 
 
-def test_merger_tilt(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger_tilt(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad", tilt=1.)
@@ -331,7 +346,7 @@ def test_merger_tilt(lattice, tws0, method, parametr=None, update_ref_values=Fal
     assert check_result(result + result2)
 
 
-def test_merger_write_read(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_merger_write_read(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
 
     R = lattice_transfer_map(lattice, energy=tws0.E)
@@ -348,7 +363,7 @@ def test_merger_write_read(lattice, tws0, method, parametr=None, update_ref_valu
     result2 = check_matrix(lattice.T, new_lat2.T, tolerance=1.0e-8, tolerance_type='absolute', assert_info=' t_matrix - ')
     assert check_result(result + result2)
 
-def test_matrix_write_read(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_matrix_write_read(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     m = Matrix(l=0.3, delta_e=0.1)
     m.r = np.random.random((6, 6))
@@ -368,11 +383,11 @@ def test_matrix_write_read(lattice, tws0, method, parametr=None, update_ref_valu
 
     result = check_matrix(R, R2, TOL, assert_info='r_matrix - ')
     result2 = check_matrix(lat.T, lat2.T, TOL, assert_info='t_matrix - ')
-    result3 = check_matrix(np.array([lat.E, lat.totalLen]), np.array([lat2.E,lat2.totalLen ]), TOL, assert_info='t_matrix - ')
+    result3 = check_matrix(np.array([ lat.totalLen]), np.array([lat2.totalLen ]), TOL, assert_info='t_matrix - ')
     assert check_result(result + result2 + result3)
 
 
-def test_matrix_b_vector(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_matrix_b_vector(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad")
@@ -407,7 +422,7 @@ def test_matrix_b_vector(lattice, tws0, method, parametr=None, update_ref_values
     assert check_result(result + result2 + result3)
 
 
-def test_matrix_b_vector_read_write(lattice, tws0, method, parametr=None, update_ref_values=False):
+def test_matrix_b_vector_read_write(lattice, tws0, method, parameter=None, update_ref_values=False):
     """R maxtrix calculation test"""
     d = Drift(l=0.5)
     q = Quadrupole(l=0.3, k1=3, k2=3.3, eid="quad")
@@ -481,16 +496,22 @@ def test_update_ref_values(lattice, tws0, method, cmdopt):
     update_functions = []
     update_functions.append('test_original_lattice_transfer_map')
     update_functions.append('test_original_twiss')
+    update_functions.append("test_lattice_transfer_maps_check")
 
     # function test_lat2input function need not be added here.
     # It is used reference results from test_original_lattice_transfer_map and test_original_twiss functions
+    update_function_parameters = {}
+    update_function_parameters['test_lattice_transfer_maps_check'] = [0, 1, 2]
+
+    parameter = update_function_parameters[cmdopt] if cmdopt in update_function_parameters.keys() else ['']
 
     if cmdopt in update_functions:
-        result = eval(cmdopt)(lattice, tws0, method, None, True)
-        if result is None:
-            return
+        for p in parameter:
+            result = eval(cmdopt)(lattice, tws0, method, p, True)
+            if result is None:
+                return
 
-        if os.path.isfile(REF_RES_DIR + cmdopt + '.json'):
-            os.rename(REF_RES_DIR + cmdopt + '.json', REF_RES_DIR + cmdopt + '.old')
+            if os.path.isfile(REF_RES_DIR + cmdopt + '.json'):
+                os.rename(REF_RES_DIR + cmdopt + '.json', REF_RES_DIR + cmdopt + str(p) + '.old')
 
-        json_save(result, REF_RES_DIR + cmdopt + '.json')
+            json_save(result, REF_RES_DIR + cmdopt + str(p) + '.json')

unit_tests/ebeam_test/io_lattice/ref_results/test_lattice_transfer_maps_check0.json

@@ -0,0 +1 @@
+{"0": 0.0, "1": 0.0, "2": 0.0, "3": 0.0, "4": 0.0, "5": 0.0}

unit_tests/ebeam_test/io_lattice/ref_results/test_lattice_transfer_maps_check1.json

@@ -0,0 +1 @@
+{"0": -0.3669735365228965, "1": -1.1490635074188513, "2": 1.190282708588602e-10, "3": 3.1482812735754887e-10, "4": 1.132202512078879e-22, "5": 2.8128820894915816e-09, "6": 0.1347434520651659, "7": 0.3173235504195618, "8": 2.2852523065905037e-11, "9": 7.028882971223359e-11, "10": 6.818683204158491e-23, "11": 1.6940566422179824e-09, "12": 2.6540664561105906e-10, "13": 7.45020103770362e-10, "14": 1.0266211358897732, "15": 2.4501975772363154, "16": 3.5480318787650113e-22, "17": 8.81484999284599e-09, "18": -9.44327601378983e-12, "19": -2.209322846125279e-11, "20": 0.055499305735640936, "21": 0.1698423074184009, "22": 6.328982779931208e-23, "23": 1.5723938178315115e-09, "24": -2.6210674684298977e-08, "25": -7.43652080033237e-08, "26": 2.9467543946615216e-08, "27": 6.169766289532141e-08, "28": 1.0, "29": -0.04189680038806823, "30": 0.0, "31": 0.0, "32": 0.0, "33": 0.0, "34": 1.5367193273535444e-15, "35": 0.03817877302615135}