Skip to content

Commit

Permalink
Added some description about monopole, dipole and quadrupole componen…
Browse files Browse the repository at this point in the history 
…ts to wake theory in the wake tutorial

added to the MagneticLattice.transfer_maps(output_at_each_step=False)
  • Loading branch information
@sergey-tomin
sergey-tomin committed Nov 9, 2022
1 parent 81782a9 commit a9edf7e
Show file tree
Hide file tree
Showing 10 changed files with 129 additions and 65 deletions.
58 changes: 57 additions & 1 deletion demos/ipython_tutorials/4_wake.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -511,7 +511,63 @@
"For some FEL applications, e.g. a two-color scheme, only one flat corrugated structure can be used to get a correlated transverse kick along the electron bunch. In that case, we can use analytical approach from [I. Zagorodnov, G. Feng, T. Limberg. Corrugated structure insertion for extending the SASE bandwidth up to 3% at the European XFEL](https://arxiv.org/abs/1607.07642) and [K. Bane, G. Stupakov, and I. Zagorodnov, Wakefields of a Beam near a Single Plate in a Flat Dechirper](https://www.slac.stanford.edu/cgi-wrap/getdoc/slac-pub-16881.pdf) to calculate described above the wakefield tables. \n",
"\n",
"<div class=\"alert alert-block alert-warning\">\n",
"<b>Note:</b> Due to the use of assumptions in the analytical approach, a transverse kick is infinite if the electron beam distance to the plate wall is zero. </div>"
"<b>Note:</b> Due to the use of assumptions in the analytical approach, a transverse kick is infinite if the electron beam distance to the plate wall is zero. </div>\n",
"\n",
"#### A small note about transverse modes\n",
"Let's look at the wakefield components. Let's start with the expression above:\n",
"$$\n",
"w_z(x_s, y_s, x_o, y_o, s) = \n",
"\\begin{bmatrix}\n",
" 1 \\\\\n",
" x_s\\\\\n",
" y_s\\\\\n",
" x_o \\\\\n",
" y_o\n",
"\\end{bmatrix}^T \n",
"\\begin{bmatrix}\n",
" h_{00}(s) & h_{01}(s) & h_{02}(s) & h_{03}(s) & h_{04}(s) \\\\\n",
" 0 & h_{11}(s) & h_{12}(s) & h_{13}(s) & h_{14}(s)\\\\\n",
" 0 & h_{12}(s) & -h_{11}(s) & h_{23}(s) & h_{24}(s) \\\\\n",
" 0 & h_{13}(s) & h_{23}(s) & h_{33}(s) & h_{34}(s)\\\\\n",
" 0 & h_{14}(s) & h_{24}(s) & h_{34}(s) & -h_{33}(s)\n",
"\\end{bmatrix} \n",
"\\begin{bmatrix}\n",
" 1 \\\\\n",
" x_s\\\\\n",
" y_s\\\\\n",
" x_o \\\\\n",
" y_o\n",
"\\end{bmatrix} ;\n",
"$$\n",
"\n",
"Expand the result of the multiplication:\n",
"$$\n",
"w_z(x_s, y_s, x_o, y_o, s) = h_{00}(s) + h_{03}(s) x_o + h_{33}(s) x_o^2 + h_{01}(s) x_s + 2 h_{13}(s) x_o x_s + h_{11}(s) x_s^2 + h_{04}(s) y_o + \n",
" 2 h_{34}(s) x_o y_o + 2 h_{14}(s) x_s y_o - h_{33}(s) y_o^2 + h_{02}(s) y_s + 2 h_{23}(s) x_o y_s + \n",
" 2 h_{12}(s) x_s y_s + 2 h_{24}(s) y_o y_s - h_{11}(s) y_s^2\n",
"$$\n",
"\n",
"Using Panofsky-Wenzel-Theorem the transverse components are:\n",
"\n",
"$$\n",
"\\frac{\\partial}{\\partial s}w_y(x_s, y_s, x_o, y_o, s) = -\\frac{\\partial}{\\partial y_o}w_z(x_s, y_s, x_o, y_o, s) = -(h_{04}(s) + 2 h_{34}(s) x_o + 2 h_{14}(s) x_s - 2 h_{33}(s) y_o + 2 h_{24}(s) y_s)\n",
"$$\n",
"where $-h_{04}(s)$ monopole transverse component, $-h_{24}(s)$ - dipole component, $h_{33}(s)$ - quadrupole component, \n",
"\n",
"\n",
"$$\n",
"\\frac{\\partial}{\\partial s}w_x(x_s, y_s, x_o, y_o, s) = -\\frac{\\partial}{\\partial x_o}w_z(x_s, y_s, x_o, y_o, s) = -(h_{03}(s) + 2 h_{33}(s) x_o + 2 h_{13}(s) x_s + 2 h_{34}(s) y_o + 2 h_{23}(s) y_s)\n",
"$$\n",
"where $h_{03}(s)$ - monopole transverse component, $-h_{13}(s)$ - dipole component, $-h_{33}(s)$ - quadrupole component. \n",
"\n",
"\n",
"Integral forms:\n",
"$$\n",
"w_y(x_s, y_s, x_o, y_o, s) = - \\int_{-\\infty}^s{\\frac{\\partial}{\\partial y_o}w_z(x_s, y_s, x_o, y_o, s') ds'}\n",
"$$\n",
"$$\n",
"w_x(x_s, y_s, x_o, y_o, s) = - \\int_{-\\infty}^s{\\frac{\\partial}{\\partial x_o}w_z(x_s, y_s, x_o, y_o, s') ds'}\n",
"$$"
]
},
{
Expand Down
73 changes: 22 additions & 51 deletions demos/ipython_tutorials/small_useful_features.ipynb
View file

Large diffs are not rendered by default.

30 changes: 22 additions & 8 deletions ocelot/cpbd/magnetic_lattice.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
MarkersInsertionReturnType = Mapping[Element, Sequence[Tuple[Marker, Marker]]]



def lattice_format_converter(elements):
"""
:param elements: lattice in the format: [[elem1, center_pos1], [elem2, center_pos2], [elem3, center_pos3], ... ]
Expand All @@ -33,7 +32,7 @@ def lattice_format_converter(elements):
s_pos = 0.0
for element in elements:
element_start = element[1] - element[0].l / 2.0
if element_start < s_pos - 1.0e-14: # 1.0e-14 is used as crutch for precision of float
if element_start < s_pos - 1.0e-14: # 1.0e-14 is used as crutch for precision of float
if element[0].l == 0.0:

if s_pos - element_start > 1.0e-2:
Expand All @@ -52,9 +51,9 @@ def lattice_format_converter(elements):
print("************** ERROR! Element " + element[0].id + " has bad position (overlapping?)")
exit()

if element_start > s_pos + 1.0e-12: # 1.0e-12 is used as crutch for precision of float
if element_start > s_pos + 1.0e-12: # 1.0e-12 is used as crutch for precision of float
drift_num += 1
drift_l = round(element_start - s_pos, 10) # round() is used as crutch for precision of float
drift_l = round(element_start - s_pos, 10) # round() is used as crutch for precision of float
drift_eid = 'D_' + str(drift_num)
cell.append(Drift(l=drift_l, eid=drift_eid))

Expand Down Expand Up @@ -125,6 +124,7 @@ def flatten(iterable: Iterator[Any]) -> Generator[Any, None, None]:
:raises: RecursionError
"""

def _flatten(iterable):
try:
for item in iterable:
Expand All @@ -134,8 +134,9 @@ def _flatten(iterable):
yield item
else:
yield from _flatten(item)
except TypeError: # If iterable isn't actually iterable, then yield it.
except TypeError: # If iterable isn't actually iterable, then yield it.
yield iterable

try:
yield from _flatten(iterable)
except RecursionError:
Expand Down Expand Up @@ -207,7 +208,7 @@ def update_transfer_maps(self):
if element.field_file is not None:
element.l = element.field_map.l * element.field_map.field_file_rep
if element.field_map.units == "mm":
element.l = element.l*0.001
element.l = element.l * 0.001
self.totalLen += element.l

tm_class_type = self.method.get(element.__class__)
Expand Down Expand Up @@ -236,7 +237,8 @@ def update_endings(self, lat_index, element, body_elements, element_util):
_logger.debug("Backtracking?")
element_util.update_last(element, body)
else:
_logger.error(element.__class__.__name__ + " is not updated. Use standard function to create and update MagneticLattice")
_logger.error(
element.__class__.__name__ + " is not updated. Use standard function to create and update MagneticLattice")

def __str__(self):
line = "LATTICE: length = " + str(self.totalLen) + " m \n"
Expand Down Expand Up @@ -279,21 +281,28 @@ 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):
def transfer_maps(self, energy, output_at_each_step: bool = False):
"""
Function calculates transfer maps, the first and second orders (R, T), for the whole lattice.
:param energy: the initial electron beam energy [GeV]
:param output_at_each_step: return three list of matrices [Bs], [Rs], [Ts] after each element in the line
:return: B, R, T - matrices
"""
Ra = np.eye(6)
Ta = np.zeros((6, 6, 6))
Ba = np.zeros((6, 1))
Bs, Rs, Ts = [], [], []
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()
Bs.append(Ba)
Rs.append(Ra)
Ts.append(Ta)
if output_at_each_step:
return Bs, Rs, Ts
return Ba, Ra, Ta


Expand Down Expand Up @@ -360,6 +369,7 @@ def exclude_zero_length_element(cell, elem_type=[UnknownElement, Marker], except
print("Exclude elements -> Element numbers: before -> after: ", len(cell), "->", len(new_cell))
return new_cell


def insert_markers_by_name(sequence, string: str, regex=False,
before=True, after=True) -> MarkersInsertionReturnType:
"""Insert markers either side of elements in the magnetic lattice, selected
Expand All @@ -380,12 +390,15 @@ def insert_markers_by_name(sequence, string: str, regex=False,
if regex:
def fre(ele):
return bool(re.match(string, ele.id))

return insert_markers_by_predicate(sequence, fre)

def f(ele):
return ele.id == string

return insert_markers_by_predicate(sequence, f, before=before, after=after)


def insert_markers_by_type(sequence, magnet_type: Element, before=True, after=True
) -> MarkersInsertionReturnType:
"""Insert markers either side of elements in the magnetic lattice, selected
Expand All @@ -405,6 +418,7 @@ def f(ele):

return insert_markers_by_predicate(sequence, f, before=before, after=after)


def insert_markers_by_predicate(sequence, predicate: Callable[[Element], bool],
before=True, after=True
) -> MarkersInsertionReturnType:
Expand Down
27 changes: 22 additions & 5 deletions unit_tests/ebeam_test/dogleg/dogleg_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,20 +13,34 @@

np.random.seed(1)

def test_lattice_transfer_map(lattice, parameter=None, update_ref_values=False):
@pytest.mark.parametrize('parameter', [0, 1, 2])
def test_lattice_transfer_map(lattice, parameter, update_ref_values=False):
"""R matrix calculation test"""

r_matrix = lattice_transfer_map(lattice, 0.0)
matrices = lattice.transfer_maps(0.0)

if update_ref_values:
return numpy2json(r_matrix)
return numpyBRT2json(matrices[parameter])

r_matrix_ref = json2numpy(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + '.json'))
matrix_ref = json2numpyBRT(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + str(parameter) + '.json'))

result = check_matrix(r_matrix, r_matrix_ref, TOL, assert_info=' r_matrix - ')
result = check_matrix(matrices[parameter], matrix_ref, TOL, assert_info=' r_matrix - ')
assert check_result(result)


@pytest.mark.parametrize('parameter', [0, 1, 2])
def test_lattice_transfer_map_at_each_step(lattice, parameter, update_ref_values=False):
"""R matrix calculation test"""

matrices = lattice.transfer_maps(0.0, output_at_each_step=True)

if update_ref_values:
return numpyBRT2json(matrices[parameter][111])

matrix_ref = json2numpyBRT(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + str(parameter) + '.json'))

result = check_matrix(matrices[parameter][111], matrix_ref, TOL, assert_info=' r_matrix - ')
assert check_result(result)


@pytest.mark.parametrize('parameter', [0, 1, 2])
Expand Down Expand Up @@ -160,10 +174,13 @@ def test_update_ref_values(lattice, cmdopt):

update_functions = []
update_functions.append('test_lattice_transfer_map')
update_functions.append('test_lattice_transfer_map_at_each_step')
update_functions.append('test_track_with_energy_shift')
update_functions.append('test_track_with_energy_shift_tilted')

update_function_parameters = {}
update_function_parameters['test_lattice_transfer_map'] = [0, 1, 2]
update_function_parameters['test_lattice_transfer_map_at_each_step'] = [0, 1, 2]
update_function_parameters['test_track_with_energy_shift'] = [0, 1, 2]
update_function_parameters['test_track_with_energy_shift_tilted'] = [0, 1, 2]
#update_functions.append('test_track_undulator_with_csr')
Expand Down
1 change: 1 addition & 0 deletions unit_tests/ebeam_test/dogleg/ref_results/test_lattice_transfer_map0.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"0": 0.0, "1": 0.0, "2": 0.0, "3": 0.0, "4": 0.0, "5": 0.0}
1 change: 1 addition & 0 deletions unit_tests/ebeam_test/dogleg/ref_results/test_lattice_transfer_map1.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"0": 0.8747331086618696, "1": -0.22264292824428675, "2": -1.9746561137456662e-11, "3": 5.682470120990124e-11, "4": 0.0, "5": -4.0494547444066366e-17, "6": 1.0547923222353865, "7": 0.874733036981342, "8": -1.1729968516195335e-10, "9": 4.4279638182207473e-11, "10": 0.0, "11": -1.5002181312074977e-17, "12": 4.427965850005827e-11, "13": 5.682470258209158e-11, "14": 1.1341829790943758, "15": -0.24736496151723855, "16": 0.0, "17": 1.2480263469995995e-07, "18": -1.1729972515479875e-10, "19": -1.9746607035085272e-11, "20": -1.1576863122040792, "21": 1.1341829790943783, "22": 0.0, "23": 4.788215180742243e-08, "24": 4.6349953123838365e-17, "25": 4.394736556336416e-17, "26": 1.9878942353335205e-07, "27": -1.5339339078102032e-07, "28": 1.0, "29": -0.029701544203881065, "30": 0.0, "31": 0.0, "32": 0.0, "33": 0.0, "34": 0.0, "35": 1.0}
1 change: 1 addition & 0 deletions unit_tests/ebeam_test/dogleg/ref_results/test_lattice_transfer_map2.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"0": -2.672042528923975e-12, "1": -3.994619303757262e-12, "2": -7.93124424756259e-06, "3": -2.6620553406949904e-05, "4": 0.0, "5": -22.971770939314915, "6": -3.994619303757437e-12, "7": -7.107080364946998e-12, "8": 2.456167262366116e-05, "9": -4.888127455773631e-05, "10": 0.0, "11": -40.266879720598865, "12": -7.931244272731585e-06, "13": 2.456167269722214e-05, "14": 2.260397095979259e-11, "15": -3.672936667049092e-12, "16": 0.0, "17": 3.5771107082825778e-09, "18": -2.6620553407329662e-05, "19": -4.888127457889589e-05, "20": -3.672936667043043e-12, "21": 1.7018256524920851e-12, "22": 0.0, "23": -3.1854234111344162e-09, "24": 0.0, "25": 0.0, "26": 0.0, "27": 0.0, "28": 0.0, "29": 0.0, "30": -22.971770939314915, "31": -40.266879720598865, "32": 3.5771107082825778e-09, "33": -3.1854234111344146e-09, "34": 0.0, "35": 2.906962596214749e-13, "36": -1.0522049609076369e-12, "37": -1.8727840503774403e-12, "38": 1.1066586387470378e-05, "39": -1.9601632363364097e-05, "40": 0.0, "41": -10.262441076443977, "42": -1.8727840503776293e-12, "43": -4.091809433635496e-12, "44": 3.473206580696664e-05, "45": -2.7333538072993324e-05, "46": 0.0, "47": -22.97177102353209, "48": 1.1066586373426057e-05, "49": 3.473206584390931e-05, "50": 8.685079828746007e-12, "51": 1.885605381581527e-12, "52": 0.0, "53": 3.265935002670562e-09, "54": -1.960163236441881e-05, "55": -2.7333538076185215e-05, "56": 1.885605381582729e-12, "57": 9.818875991366078e-14, "58": 0.0, "59": -1.887560090462635e-09, "60": 0.0, "61": 0.0, "62": 0.0, "63": 0.0, "64": 0.0, "65": 0.0, "66": -10.262441076443977, "67": -22.97177102353209, "68": 3.265935002670562e-09, "69": -1.8875600904626337e-09, "70": 0.0, "71": 1.6360065546611907e-13, "72": -1.6859997357071716e-05, "73": -3.246136273432771e-05, "74": -2.2491783244960533e-12, "75": 5.760119416340184e-13, "76": 0.0, "77": -1.8876001350588844e-09, "78": -3.246136273529544e-05, "79": -4.815442476183793e-05, "80": 2.2000397786079492e-12, "81": 1.0213266125146733e-12, "82": 0.0, "83": -3.185465605199176e-09, "84": -2.249178324495978e-12, "85": 2.2000397786425147e-12, "86": -6.915522386745788e-05, "87": 2.5340043865857342e-05, "88": 0.0, "89": -0.218605541267886, "90": 5.760119416291821e-13, "91": 1.0213266124926011e-12, "92": 2.5340043874526013e-05, "93": -4.649418720485507e-05, "94": 0.0, "95": 0.5781130875405778, "96": 0.0, "97": 0.0, "98": 0.0, "99": 0.0, "100": 0.0, "101": 0.0, "102": -1.887600135058885e-09, "103": -3.185465605199176e-09, "104": -0.2186055412678865, "105": 0.5781130875405772, "106": 0.0, "107": -1.2121722199462648e-07, "108": 3.313776673817418e-05, "109": 3.709299855424934e-05, "110": -1.2034650600688e-11, "111": 4.289257414462862e-12, "112": 0.0, "113": 3.2661379038088534e-09, "114": 3.709299855358505e-05, "115": 2.338867720035043e-05, "116": -2.1387740944616892e-11, "117": 1.4275023927138033e-12, "118": 0.0, "119": 3.5770851955093723e-09, "120": -1.2034650600687729e-11, "121": -2.1387740944637556e-11, "122": -2.6809278282567163e-05, "123": 4.913716122438051e-05, "124": 0.0, "125": -0.7010019607179931, "126": 4.289257414455605e-12, "127": 1.427502392714889e-12, "128": 4.913716122260372e-05, "129": 1.1400828544399274e-05, "130": 0.0, "131": -0.21859929644815293, "132": 0.0, "133": 0.0, "134": 0.0, "135": 0.0, "136": 0.0, "137": 0.0, "138": 3.266137903808854e-09, "139": 3.5770851955093723e-09, "140": -0.7010019607179925, "141": -0.21859929644815337, "142": 0.0, "143": -8.132932037890643e-08, "144": 15.25354321563195, "145": 22.378940031588673, "146": -9.729468531236475e-10, "147": 1.7669704365423119e-09, "148": 0.0, "149": -1.641044741415737e-13, "150": 22.378940031588677, "151": 40.337207209159104, "152": -3.900317160826793e-09, "153": 3.205638975567149e-09, "154": 0.0, "155": -2.9163912941835087e-13, "156": -9.729468531236466e-10, "157": -3.900317160826793e-09, "158": -1.048141135055912, "159": 0.421342007073367, "160": 0.0, "161": -1.5555426183837641e-07, "162": 1.7669704365423127e-09, "163": 3.205638975567148e-09, "164": 0.4213420070733681, "165": -0.6016122173302836, "166": 0.0, "167": 2.125636010360868e-07, "168": 0.0, "169": 0.0, "170": 0.0, "171": 0.0, "172": 0.0, "173": 0.0, "174": -1.6410447414144444e-13, "175": -2.9163912941835087e-13, "176": -1.555542618297028e-07, "177": 2.1256360092679923e-07, "178": 0.0, "179": 0.09691709513392965, "180": 0.0, "181": 0.0, "182": 0.0, "183": 0.0, "184": 0.0, "185": 0.0, "186": 0.0, "187": 0.0, "188": 0.0, "189": 0.0, "190": 0.0, "191": 0.0, "192": 0.0, "193": 0.0, "194": 0.0, "195": 0.0, "196": 0.0, "197": 0.0, "198": 0.0, "199": 0.0, "200": 0.0, "201": 0.0, "202": 0.0, "203": 0.0, "204": 0.0, "205": 0.0, "206": 0.0, "207": 0.0, "208": 0.0, "209": 0.0, "210": 0.0, "211": 0.0, "212": 0.0, "213": 0.0, "214": 0.0, "215": 0.0}
1 change: 1 addition & 0 deletions unit_tests/ebeam_test/dogleg/ref_results/test_lattice_transfer_map_at_each_step0.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"0": 0.0, "1": 0.0, "2": 0.0, "3": 0.0, "4": 0.0, "5": 0.0}
1 change: 1 addition & 0 deletions unit_tests/ebeam_test/dogleg/ref_results/test_lattice_transfer_map_at_each_step1.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
{"0": 2.0858193324342866, "1": 3.6130148951058696, "2": -2.3138711477042311e-10, "3": 1.2606522012903262e-11, "4": 0.0, "5": -6.465587810225149e-11, "6": 1.0995593365631242, "7": 2.384062791887065, "8": -1.4812129349624275e-10, "9": -1.5087106433316238e-12, "10": 0.0, "11": -3.86723800597904e-11, "12": -4.8562109650334193e-11, "13": 7.581851159690206e-12, "14": -0.17447969216250495, "15": 0.8361294426483068, "16": 0.0, "17": 0.07723998829085041, "18": -1.7806686836051348e-12, "19": -4.536231643892462e-11, "20": -1.129297842125262, "21": -0.31958346549542405, "22": 0.0, "23": -0.042588821587950514, "24": -7.364881676834986e-12, "25": 1.7600670759387554e-11, "26": 0.09465783658286767, "27": -0.010925144524565361, "28": 1.0, "29": -0.018566904150905235, "30": 0.0, "31": 0.0, "32": 0.0, "33": 0.0, "34": 0.0, "35": 1.0}
Loading

0 comments on commit a9edf7e

Please sign in to comment.