Skip to content

Commit 5a480f2

Browse files
i-a-morozovi-a-morozov
committed
12_03_2024: added frequency from trajectory example (linear system and parametric derivatives)
1 parent 3abc0ea commit 5a480f2

File tree

2 files changed

+299
-0
lines changed

2 files changed

+299
-0
lines changed

‎docs/source/examples/model-57.ipynb‎

Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
{
2+
"cells": [
3+
{
4+
"attachments": {},
5+
"cell_type": "markdown",
6+
"id": "262a5ec8-2553-4237-ab62-319b6ca22089",
7+
"metadata": {},
8+
"source": [
9+
"# Example-58: Frequency (parametric derivatives for linear system)"
10+
]
11+
},
12+
{
13+
"cell_type": "code",
14+
"execution_count": 1,
15+
"id": "21defbfe-5f6e-4e91-94dd-a8267d6144d1",
16+
"metadata": {},
17+
"outputs": [],
18+
"source": [
19+
"# In this example frequencies for linear model are computed from trajectory data\n",
20+
"# Frequency values are compared with ones obtained from one-turn matrix\n",
21+
"# Derivatives with respect to parameters are also computed and compared"
22+
]
23+
},
24+
{
25+
"cell_type": "code",
26+
"execution_count": 2,
27+
"id": "6302e8bb-59b6-47fc-a81d-717dadda6e06",
28+
"metadata": {},
29+
"outputs": [],
30+
"source": [
31+
"# Import\n",
32+
"\n",
33+
"import torch\n",
34+
"from torch import Tensor\n",
35+
"\n",
36+
"from pathlib import Path\n",
37+
"\n",
38+
"from model.library.line import Line\n",
39+
"\n",
40+
"from model.command.external import load_sdds\n",
41+
"from model.command.external import load_lattice\n",
42+
"from model.command.build import build\n",
43+
"from model.command.tune import tune\n",
44+
"from model.command.tune import chromaticity\n",
45+
"from model.command.trajectory import trajectory\n",
46+
"from model.command.frequency import filter\n",
47+
"from model.command.frequency import frequency_factory"
48+
]
49+
},
50+
{
51+
"cell_type": "code",
52+
"execution_count": 3,
53+
"id": "93add4df-b6ac-475f-b641-d00cf202ea3a",
54+
"metadata": {},
55+
"outputs": [],
56+
"source": [
57+
"# Load ELEGANT twiss\n",
58+
"\n",
59+
"path = Path('ic.twiss')\n",
60+
"parameters, columns = load_sdds(path)\n",
61+
"\n",
62+
"NUX:Tensor = torch.tensor(parameters['nux'] % 1, dtype=torch.float64)\n",
63+
"NUY:Tensor = torch.tensor(parameters['nuy'] % 1, dtype=torch.float64)"
64+
]
65+
},
66+
{
67+
"cell_type": "code",
68+
"execution_count": 4,
69+
"id": "4aa78283-6227-47bd-a8d0-816e2bf3bc03",
70+
"metadata": {},
71+
"outputs": [],
72+
"source": [
73+
"# Build and setup lattice\n",
74+
"\n",
75+
"# Load ELEGANT table\n",
76+
"\n",
77+
"path = Path('ic.lte')\n",
78+
"data = load_lattice(path)\n",
79+
"\n",
80+
"# Build ELEGANT table\n",
81+
"\n",
82+
"ring:Line = build('RING', 'ELEGANT', data)\n",
83+
"ring.flatten()\n",
84+
"\n",
85+
"# Merge drifts\n",
86+
"\n",
87+
"ring.merge()\n",
88+
"\n",
89+
"# Turn off sextupoles and set linear dipoles\n",
90+
"\n",
91+
"for element in ring: \n",
92+
" if element.__class__.__name__ == 'Sextupole':\n",
93+
" element.ms = 0.0\n",
94+
" if element.__class__.__name__ == 'Dipole':\n",
95+
" element.linear = True \n",
96+
"\n",
97+
"# Split BPMs\n",
98+
"\n",
99+
"ring.split((None, ['BPM'], None, None))\n",
100+
"\n",
101+
"# Roll lattice start\n",
102+
"\n",
103+
"ring.roll(1)\n",
104+
"\n",
105+
"# Split lattice into lines by BPMs\n",
106+
"\n",
107+
"ring.splice()\n",
108+
"\n",
109+
"# Set number of elements of different kinds\n",
110+
"\n",
111+
"nb = ring.describe['BPM']\n",
112+
"nq = ring.describe['Quadrupole']\n",
113+
"ns = ring.describe['Sextupole']"
114+
]
115+
},
116+
{
117+
"cell_type": "code",
118+
"execution_count": 5,
119+
"id": "59f8426b-30da-4540-b98c-d43720b3e0e8",
120+
"metadata": {},
121+
"outputs": [
122+
{
123+
"name": "stdout",
124+
"output_type": "stream",
125+
"text": [
126+
"tensor(3.1086e-15, dtype=torch.float64)\n",
127+
"tensor(5.5511e-16, dtype=torch.float64)\n"
128+
]
129+
}
130+
],
131+
"source": [
132+
"# Compute tunes (one-turn matrix)\n",
133+
"\n",
134+
"nux, nuy = tune(ring, [], alignment=False, matched=True, limit=8, epsilon=1.0E-12)\n",
135+
"\n",
136+
"# Compare with elegant\n",
137+
"\n",
138+
"print((NUX - nux).abs())\n",
139+
"print((NUY - nuy).abs())"
140+
]
141+
},
142+
{
143+
"cell_type": "code",
144+
"execution_count": 6,
145+
"id": "bef63178-4ba8-47e2-81dd-5cc6b1f5e926",
146+
"metadata": {},
147+
"outputs": [
148+
{
149+
"name": "stdout",
150+
"output_type": "stream",
151+
"text": [
152+
"tensor(3.3061e-12, dtype=torch.float64)\n",
153+
"tensor(3.1350e-10, dtype=torch.float64)\n"
154+
]
155+
}
156+
],
157+
"source": [
158+
"# Compute tunes (trajectory)\n",
159+
"\n",
160+
"# Set trajectory generator\n",
161+
"\n",
162+
"generator = trajectory(ring, [0], matched=True)\n",
163+
"\n",
164+
"# Set initial condition\n",
165+
"\n",
166+
"state = torch.tensor([+1.0E-9, 0.0, -1.0E-9, 0.0], dtype=torch.float64)\n",
167+
"\n",
168+
"# Set window data\n",
169+
"\n",
170+
"window = filter(2**10, 1.0, dtype=ring.dtype, device=ring.device)\n",
171+
"\n",
172+
"# Set frequency generator\n",
173+
"\n",
174+
"frequency = frequency_factory(generator)\n",
175+
"\n",
176+
"# Compute frequencies\n",
177+
"\n",
178+
"nux, nuy = frequency(window, state)\n",
179+
"\n",
180+
"# Compare with elegant\n",
181+
"\n",
182+
"print((NUX - nux).abs())\n",
183+
"print((NUY - nuy).abs())"
184+
]
185+
},
186+
{
187+
"cell_type": "code",
188+
"execution_count": 7,
189+
"id": "43d44857-d79a-469d-8f38-b167ad480efa",
190+
"metadata": {},
191+
"outputs": [
192+
{
193+
"name": "stdout",
194+
"output_type": "stream",
195+
"text": [
196+
"tensor([[-7.8819],\n",
197+
" [-3.9483]], dtype=torch.float64)\n"
198+
]
199+
}
200+
],
201+
"source": [
202+
"# Derivative with respect to momentum deviation (one-turn matrix)\n",
203+
"\n",
204+
"dp = torch.tensor([0.0], dtype=torch.float64)\n",
205+
"\n",
206+
"print(torch.func.jacrev(lambda dp: tune(ring, [dp], ('dp', None, None, None), matched=True, limit=1, epsilon=None))(dp))"
207+
]
208+
},
209+
{
210+
"cell_type": "code",
211+
"execution_count": 8,
212+
"id": "827913f7-2278-40c2-b88d-d41758c3192d",
213+
"metadata": {},
214+
"outputs": [
215+
{
216+
"name": "stdout",
217+
"output_type": "stream",
218+
"text": [
219+
"tensor([[-7.8819],\n",
220+
" [-3.9486]], dtype=torch.float64)\n"
221+
]
222+
}
223+
],
224+
"source": [
225+
"# Derivative with respect to momentum deviation (trajectory)\n",
226+
"\n",
227+
"# Set parametric trajectory generator\n",
228+
"\n",
229+
"generator = trajectory(ring, [0], ('dp', None, None, None), matched=True)\n",
230+
"\n",
231+
"# Set initial state and momentum deviation\n",
232+
"# Note, state should not be equal to zero, since zero is a fixed point,\n",
233+
"\n",
234+
"state = torch.tensor([+1.0E-9, 0.0, -1.0E-9, 0.0], dtype=torch.float64)\n",
235+
"dp = torch.tensor([0.0], dtype=torch.float64)\n",
236+
"\n",
237+
"# Set window data\n",
238+
"\n",
239+
"window = filter(2**10, 1.0, dtype=ring.dtype, device=ring.device)\n",
240+
"\n",
241+
"# Set frequency generator\n",
242+
"\n",
243+
"frequency = frequency_factory(generator)\n",
244+
"\n",
245+
"# Compute derivative\n",
246+
"\n",
247+
"print(torch.func.jacrev(lambda dp: frequency(window, state, dp), chunk_size=256)(dp))"
248+
]
249+
}
250+
],
251+
"metadata": {
252+
"colab": {
253+
"collapsed_sections": [
254+
"myt0_gMIOq7b",
255+
"5d97819c"
256+
],
257+
"name": "03_frequency.ipynb",
258+
"provenance": []
259+
},
260+
"kernelspec": {
261+
"display_name": "Python 3 (ipykernel)",
262+
"language": "python",
263+
"name": "python3"
264+
},
265+
"language_info": {
266+
"codemirror_mode": {
267+
"name": "ipython",
268+
"version": 3
269+
},
270+
"file_extension": ".py",
271+
"mimetype": "text/x-python",
272+
"name": "python",
273+
"nbconvert_exporter": "python",
274+
"pygments_lexer": "ipython3",
275+
"version": "3.12.1"
276+
},
277+
"latex_envs": {
278+
"LaTeX_envs_menu_present": true,
279+
"autoclose": false,
280+
"autocomplete": true,
281+
"bibliofile": "biblio.bib",
282+
"cite_by": "apalike",
283+
"current_citInitial": 1,
284+
"eqLabelWithNumbers": true,
285+
"eqNumInitial": 1,
286+
"hotkeys": {
287+
"equation": "Ctrl-E",
288+
"itemize": "Ctrl-I"
289+
},
290+
"labels_anchors": false,
291+
"latex_user_defs": false,
292+
"report_style_numbering": false,
293+
"user_envs_cfg": false
294+
}
295+
},
296+
"nbformat": 4,
297+
"nbformat_minor": 5
298+
}

‎docs/source/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ Simple accelerator lattice model: linear optics errors, closed orbit and Twiss p
6565
examples/model-54.ipynb
6666
examples/model-55.ipynb
6767
examples/model-56.ipynb
68+
examples/model-57.ipynb
6869

6970
.. toctree::
7071
:caption: API:

0 commit comments

Comments
 (0)