Skip to content
Snippets Groups Projects
Unverified Commit 07f43eec authored by tegenolf's avatar tegenolf Committed by GitHub
Browse files

Add files via upload

parent aa6aa372
No related branches found
No related tags found
No related merge requests found
Show whitespace changes
Inline Side-by-side
Showing
with 889 additions and 0 deletions
week2/nmap2.ipynb 0 → 100644
+ 889
− 0
View file @ 07f43eec
%% Cell type:markdown id:guided-sense tags:
<h1 style="text-align: center; vertical-align: middle;">Numerical Methods in Accelerator Physics</h1>
<h2 style="text-align: center; vertical-align: middle;">Python examples -- Week 2</h2>
%% Cell type:markdown id:instrumental-stylus tags:
<h2>Run this first!</h2>
Imports and modules:
%% Cell type:code id:universal-freedom tags:
``` python
from config2 import *
%matplotlib inline
```
%% Cell type:markdown id:conventional-header tags:
<h3>Pendulum parameters, Hamiltonian</h3>
%% Cell type:code id:cellular-passport tags:
``` python
m = 1 # point mass
g = 1 # magnitude of the gravitational field
L = 1 # length of the rod
```
%% Cell type:code id:worthy-bridges tags:
``` python
### Values of hamiltonian for comparison with theory
TH, PP = np.meshgrid(np.linspace(-np.pi * 1.1, np.pi * 1.1, 100),
np.linspace(-3, 3, 100))
HH = hamiltonian(TH, PP)
```
%% Cell type:markdown id:reverse-algebra tags:
<h3>Euler-Cromer Method (slide 10)</h3>
%% Cell type:code id:broke-edward tags:
``` python
def solve_eulercromer(theta, p, dt=0.1):
### from Euler:
theta_next = theta + dt * p / (m * L*L)
p_next = p - dt * m * g * L * np.sin(theta_next)
return (theta_next, p_next)
```
%% Cell type:code id:pediatric-prompt tags:
``` python
### Initial values:
theta_ini = -1.1
p_ini = 0
n_steps = 100
```
%% Cell type:code id:italic-muscle tags:
``` python
results_eulercromer = np.zeros((n_steps, 2), dtype=np.float32)
results_eulercromer[0] = (theta_ini, p_ini)
for k in range(1, n_steps):
results_eulercromer[k] = solve_eulercromer(*results_eulercromer[k - 1])
```
%% Cell type:code id:leading-liberal tags:
``` python
### Explicit Euler for comparison:
results_euler = np.zeros((n_steps, 2), dtype=np.float32) # for comparison
results_euler[0] = (theta_ini, p_ini)
for k in range(1, n_steps):
results_euler[k] = solve_euler(*results_euler[k - 1])
```
%% Cell type:code id:posted-irrigation tags:
``` python
plt.contour(TH, PP, HH, levels=10, linestyles='--', linewidths=1)
plt.plot(results_euler[:, 0], results_euler[:, 1], c='b', label='Euler')
plt.plot(results_eulercromer[:, 0], results_eulercromer[:, 1], c='c', label='Euler-Cromer')
plt.legend(bbox_to_anchor=(1.05, 1))
set_axes()
```
%% Cell type:code id:scientific-comparison tags:
``` python
plt.plot(
hamiltonian(results_euler[:, 0], results_euler[:, 1]),
c='b', label='Euler')
plt.plot(
hamiltonian(results_eulercromer[:, 0], results_eulercromer[:, 1]),
c='c', label='Euler-Cromer')
plt.axhline(hamiltonian(theta_ini, p_ini), c='r', label='theory')
plt.xlabel('Steps $k$')
plt.ylabel(r'$\mathcal{H}(\theta, p)$')
plt.legend(bbox_to_anchor=(1.05, 1));
```
%% Cell type:markdown id:fitted-occasions tags:
<h3>Leapfrog Method (slide 12)</h3>
%% Cell type:code id:lovely-silver tags:
``` python
def solve_leapfrog(theta, p, dt=0.1):
theta_half = theta + dt / 2 * p / (m * L*L)
p_next = p - dt * m * g * L * np.sin(theta_half)
theta_next = theta_half + dt / 2 * p_next / (m * L*L)
return (theta_next, p_next)
```
%% Cell type:code id:civil-tunnel tags:
``` python
results_leapfrog = np.zeros((n_steps, 2), dtype=np.float32)
results_leapfrog[0] = (theta_ini, p_ini)
for k in range(1, n_steps):
results_leapfrog[k] = solve_leapfrog(*results_leapfrog[k - 1])
```
%% Cell type:code id:determined-consortium tags:
``` python
plt.plot(
hamiltonian(results_euler[:, 0], results_euler[:, 1]),
c='b', label='Euler, $\mathcal{O}(\Delta t^2)$')
plt.plot(
hamiltonian(results_eulercromer[:, 0], results_eulercromer[:, 1]),
c='c', label='Euler-Cromer, $\mathcal{O}(\Delta t^2)$')
plt.plot(
hamiltonian(results_leapfrog[:, 0], results_leapfrog[:, 1]),
c='k', label='leapfrog, $\mathcal{O}(\Delta t^3)$')
plt.axhline(hamiltonian(theta_ini, p_ini), c='r', lw=2, label='theory')
plt.ylim(0.5, 0.6)
plt.xlabel('Steps $k$')
plt.ylabel(r'$\mathcal{H}(\theta, p)$')
plt.legend(bbox_to_anchor=(1.05, 1));
```
%% Cell type:markdown id:worst-bikini tags:
<h3>Collection of pendulums (slides 14 & 17)</h3>
%% Cell type:code id:respiratory-chaos tags:
``` python
theta_grid = np.linspace(-0.5 * np.pi, 0.5 * np.pi, 21)
p_grid = np.linspace(-0.3, 0.3, 5)
thetas, ps = np.meshgrid(theta_grid, p_grid)
```
%% Cell type:code id:smart-school tags:
``` python
N = len(theta_grid) * len(p_grid)
N
```
%% Cell type:code id:australian-baker tags:
``` python
plt.scatter(thetas, ps, c='b', marker='.')
plot_hamiltonian()
```
%% Cell type:markdown id:electrical-union tags:
<h3>Time evolution</h3>
%% Cell type:code id:sixth-plain tags:
``` python
n_steps = 100
```
%% Cell type:code id:weird-lawyer tags:
``` python
results_thetas = np.zeros((n_steps, N), dtype=np.float32)
results_thetas[0] = thetas.flatten()
results_ps = np.zeros((n_steps, N), dtype=np.float32)
results_ps[0] = ps.flatten()
for k in range(1, n_steps):
results_thetas[k], results_ps[k] = solve_leapfrog(results_thetas[k - 1], results_ps[k - 1])
```
%% Cell type:markdown id:special-canadian tags:
<h3>Observations of centroids (slide 19)</h3>
%% Cell type:code id:indonesian-redhead tags:
``` python
centroids_theta = 1/N * np.sum(results_thetas, axis=1)
centroids_p = 1/N * np.sum(results_ps, axis=1)
```
%% Cell type:code id:outside-wonder tags:
``` python
plt.plot(centroids_theta, label=r'$\langle\theta\rangle$')
plt.plot(centroids_p, label=r'$\langle p\rangle$')
plt.xlabel('Steps $k$')
plt.ylabel('Centroid amplitude')
plt.legend(bbox_to_anchor=(1.05, 1));
```
%% Cell type:markdown id:forbidden-merit tags:
<h3>Observations of RMS size</h3>
%% Cell type:code id:continental-preliminary tags:
``` python
var_theta = 1/N * np.sum(results_thetas * results_thetas, axis=1)
var_p = 1/N * np.sum(results_ps * results_ps, axis=1)
```
%% Cell type:code id:metropolitan-princess tags:
``` python
plt.plot(var_theta, label=r'$\langle\theta^2\rangle$')
plt.plot(var_p, label=r'$\langle p^2\rangle$')
plt.xlabel('Steps $k$')
plt.ylabel('Variance')
plt.legend(bbox_to_anchor=(1.05, 1));
```
%% Cell type:markdown id:quick-humor tags:
<h3>Distribution evolution over time</h3>
%% Cell type:code id:homeless-deployment tags:
``` python
k = 18
plt.scatter(results_thetas[k], results_ps[k], c='b', marker='.')
plot_hamiltonian()
```
%% Cell type:markdown id:lined-glory tags:
<h3>RMS emittance evolution (slide 20) $$\epsilon = \sqrt{\langle \theta^2\rangle \langle p^2\rangle - \langle \theta\,p\rangle^2}$$</h3>
%% Cell type:code id:welsh-effort tags:
``` python
def emittance(theta, p):
N = len(theta)
# subtract centroids
theta = theta - 1/N * np.sum(theta)
p = p - 1/N * np.sum(p)
# compute Σ matrix entries
theta_sq = 1/N * np.sum(theta * theta)
p_sq = 1/N * np.sum(p * p)
crossterm = 1/N * np.sum(theta * p)
# determinant of Σ matrix
epsilon = np.sqrt(theta_sq * p_sq - crossterm * crossterm)
return epsilon
```
%% Cell type:code id:impaired-russian tags:
``` python
results_emit = np.zeros(n_steps, dtype=np.float32)
for k in range(n_steps):
results_emit[k] = emittance(results_thetas[k], results_ps[k])
```
%% Cell type:code id:enhanced-dairy tags:
``` python
plt.plot(results_emit)
plt.xlabel('Steps $k$')
plt.ylabel('RMS emittance $\epsilon$');
```
%% Cell type:markdown id:taken-survival tags:
<h3>Exercise: Linearized pendulum motion</h3>
%% Cell type:markdown id:thorough-failing tags:
$$\left.\begin{matrix}\,
\theta_{k+1/2} &=& \theta_k + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k}}{mL^2} \\
p_{k+1} &=& p_k - \Delta t\cdot m g L\sin\theta_{k+1/2} \\
\theta_{k+1} &=& \theta_{k+1/2} + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k+1}}{mL^2}
\end{matrix}\right\} \quad \stackrel{\text{Taylor}}{\Longrightarrow} \quad \left\{\begin{matrix}\,
\theta_{k+1/2} &=& \theta_k + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k}}{mL^2} \\
p_{k+1} &=& p_k - \Delta t\cdot m g L\color{red}{\left(\theta_{k+1/2}-\frac{\theta_{k+1/2}^3}{3!}+\frac{\theta_{k+1/2}^5}{5!}-\frac{\theta_{k+1/2}^7}{7!}+...\right)} \\
\theta_{k+1} &=& \theta_{k+1/2} + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k+1}}{mL^2}
\end{matrix}\right.$$
%% Cell type:markdown id:quiet-updating tags:
$$\left.\begin{matrix}\,
\theta_{k+1/2} &=& \theta_k + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k}}{mL^2} \\
p_{k+1} &=& p_k - \Delta t\cdot m g L\sin\theta_{k+1/2} \\
\theta_{k+1} &=& \theta_{k+1/2} + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k+1}}{mL^2}
\end{matrix}\right\} \quad \stackrel{\text{First order}}{\Longrightarrow} \quad \left\{\begin{matrix}\,
\theta_{k+1/2} &=& \theta_k + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k}}{mL^2} \\
p_{k+1} &=& p_k - \Delta t\cdot m g L\cdot\color{red}{\theta_{k+1/2}} \\
\theta_{k+1} &=& \theta_{k+1/2} + {\cfrac{\Delta t}{2}}\cdot\cfrac{p_{k+1}}{mL^2}
\end{matrix}\right.$$
%% Cell type:code id:living-spyware tags:
``` python
def solve_leapfrog_linear(theta, p, dt=dt):
theta_half = theta + dt / 2 * p / (m * L*L)
p_next = p - dt * m * g * L * theta_half
theta_next = theta_half + dt / 2 * p_next / (m * L*L)
return (theta_next, p_next)
```
%% Cell type:code id:bored-commodity tags:
``` python
results_thetas_linear = np.zeros((n_steps, N), dtype=np.float32)
results_thetas_linear[0] = thetas.flatten()
results_ps_linear = np.zeros((n_steps, N), dtype=np.float32)
results_ps_linear[0] = ps.flatten()
for k in range(1, n_steps):
results_thetas_linear[k], results_ps_linear[k] = solve_leapfrog_linear(results_thetas_linear[k - 1], results_ps_linear[k - 1])
```
%% Cell type:code id:quantitative-programmer tags:
``` python
var_theta_linear = 1/N * np.sum(results_thetas_linear * results_thetas_linear, axis=1)
var_p_linear = 1/N * np.sum(results_ps_linear * results_ps_linear, axis=1)
```
%% Cell type:code id:three-better tags:
``` python
plt.plot(var_theta_linear, label=r'$\langle\theta^2\rangle$')
plt.plot(var_p_linear, label=r'$\langle p^2\rangle$')
plt.xlabel('Steps $k$')
plt.ylabel('Variance')
plt.legend(bbox_to_anchor=(1.05, 1));
```
%% Cell type:code id:statutory-thriller tags:
``` python
results_emit_linear = np.zeros(n_steps, dtype=np.float32)
for k in range(n_steps):
results_emit_linear[k] = emittance(results_thetas_linear[k], results_ps_linear[k])
```
%% Cell type:code id:equipped-campus tags:
``` python
plt.plot(results_emit_linear)
plt.xlabel('Steps $k$')
plt.ylabel('RMS emittance $\epsilon$');
```
%% Cell type:markdown id:breeding-invalid tags:
<h3>Centroid offset in (original) pendulum</h3>
%% Cell type:code id:chemical-aquarium tags:
``` python
results_thetas2 = np.zeros((n_steps, N), dtype=np.float32)
results_thetas2[0] = thetas.flatten() + 0.1 * np.pi
results_ps2 = np.zeros((n_steps, N), dtype=np.float32)
results_ps2[0] = ps.flatten()
for k in range(1, n_steps):
results_thetas2[k], results_ps2[k] = solve_leapfrog(results_thetas2[k - 1], results_ps2[k - 1])
```
%% Cell type:code id:greek-brunswick tags:
``` python
centroids_theta2 = 1/N * np.sum(results_thetas2, axis=1)
centroids_p2 = 1/N * np.sum(results_ps2, axis=1)
plt.plot(centroids_theta, label=r'no offset: $\langle\theta\rangle$')
plt.plot(centroids_theta2, label=r'with offset: $\langle\theta\rangle$')
plt.plot(centroids_p2, label=r'with offset: $\langle p\rangle$')
plt.xlabel('Steps $k$')
plt.ylabel('Centroid amplitude')
plt.legend(bbox_to_anchor=(1.05, 1));
```
%% Cell type:code id:voluntary-transparency tags:
``` python
results_emit2 = np.zeros(n_steps, dtype=np.float32)
for k in range(n_steps):
results_emit2[k] = emittance(results_thetas2[k], results_ps2[k])
```
%% Cell type:code id:legitimate-objective tags:
``` python
plt.plot(results_emit, label='no offset')
plt.plot(results_emit2, label='with offset')
plt.xlabel('Steps $k$')
plt.ylabel('RMS emittance $\epsilon$')
plt.legend();
```
%% Cell type:markdown id:rough-funds tags:
<h3>Long-term behaviour</h3>
%% Cell type:code id:canadian-section tags:
``` python
n_steps = 3000
```
%% Cell type:code id:primary-lightning tags:
``` python
results_thetas3 = np.zeros((n_steps, N), dtype=np.float32)
results_thetas3[0] = thetas.flatten() + 0.1 * np.pi
results_ps3 = np.zeros((n_steps, N), dtype=np.float32)
results_ps3[0] = ps.flatten()
for k in range(1, n_steps):
results_thetas3[k], results_ps3[k] = solve_leapfrog(results_thetas3[k - 1], results_ps3[k - 1])
```
%% Cell type:code id:ordinary-advertising tags:
``` python
plot_macro_evolution(results_thetas3, results_ps3);
```
%% Cell type:markdown id:figured-attendance tags:
<h3>Employ FFT on pendulum motion (slide 31)</h3>
%% Cell type:code id:coordinated-boring tags:
``` python
## Observe two particles:
i1 = N // 2 - 10
i2 = N // 2
plt.plot(results_thetas3[:, i1])
plt.plot(results_thetas3[:, i2])
plt.xlabel('Steps $k$')
plt.ylabel(r'$\theta$');
```
%% Cell type:code id:supported-involvement tags:
``` python
## Obtain the FFT spectra for both:
spec1 = np.fft.rfft(results_thetas3[:, i1])
spec2 = np.fft.rfft(results_thetas3[:, i2])
freq = np.fft.rfftfreq(n_steps)
## Frequency of linearized system
freq_theory = np.sqrt(1 / 1) * 0.1 / (2 * np.pi)
```
%% Cell type:code id:weird-retrieval tags:
``` python
plt.plot(freq, np.abs(spec1))
plt.plot(freq, np.abs(spec2))
plt.axvline(freq_theory, c='k', zorder=0, ls='--')
plt.xlim(0, 0.05)
plt.xlabel('Phase advance per $\Delta t$ [$2\pi$]')
plt.ylabel('FFT spectrum')
plt.yscale('log');
```
%% Cell type:markdown id:desirable-venice tags:
<h3>Frequency vs Amplitude</h3>
%% Cell type:code id:textile-commons tags:
``` python
specs = np.abs(np.fft.rfft(results_thetas3.T))
```
%% Cell type:code id:standing-there tags:
``` python
max_ids = np.argmax(specs, axis=1)
amplitudes = np.max(specs, axis=1)
```
%% Cell type:code id:developed-input tags:
``` python
i_range = np.where(ps.flatten() == 0)[0]
plt.scatter(results_thetas3[0][i_range], freq[max_ids][i_range])
plt.axhline(freq_theory, c='k', ls='--', zorder=0)
plt.xticks([-np.pi/2, 0, np.pi/2, np.pi], [r"$-\pi/2$", "0", r"$\pi/2$", r"$\pi$"])
plt.xlabel(r'Initial $\theta$ ($p=0$)')
plt.ylabel('Phase advance / \n' + r'integration step $\Delta t$ [$2\pi$]');
plt.ylim(0.012, 0.0165);
```
%% Cell type:markdown id:unavailable-scenario tags:
<h3>NAFF Algorithm (slide 32)</h3>
%% Cell type:code id:present-helena tags:
``` python
import PyNAFF
```
%% Cell type:code id:global-distributor tags:
``` python
freqs_naff = []
for signal in results_thetas3.T[i_range]:
freq_naff = PyNAFF.naff(signal, turns=n_steps - 1, nterms=1)
try:
freq_naff = freq_naff[0, 1]
except IndexError:
freq_naff = 0
freqs_naff += [freq_naff]
freqs_naff = np.array(freqs_naff)
```
%% Cell type:code id:statutory-union tags:
``` python
plt.scatter(results_thetas3[0][i_range], freq[max_ids][i_range])
plt.plot(results_thetas3[0][i_range], freqs_naff, c='r', marker='.')
plt.axhline(freq_theory, c='k', ls='--', zorder=0)
plt.xticks([-np.pi/2, 0, np.pi/2, np.pi], [r"$-\pi/2$", "0", r"$\pi/2$", r"$\pi$"])
plt.xlabel(r'Initial $\theta$ ($p=0$)')
plt.ylabel('Phase advance / \n' + r'integration step $\Delta t$ [$2\pi$]');
```
%% Cell type:code id:inappropriate-listening tags:
``` python
plt.plot(results_thetas3[0][i_range], 100 * (freq[max_ids][i_range] - freqs_naff) / freqs_naff)
plt.xlabel(r'Initial $\theta$ ($p=0$)')
plt.ylabel('Error (FFT - NAFF) / NAFF ' + r'[%]');
```
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment