From 07f43eec02b8da1ff9c1efc24914d8d8267d2dc8 Mon Sep 17 00:00:00 2001 From: tegenolf <138484196+tegenolf@users.noreply.github.com> Date: Thu, 24 Oct 2024 23:41:02 +0200 Subject: [PATCH] Add files via upload --- week2/nmap2.ipynb | 889 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 889 insertions(+) create mode 100644 week2/nmap2.ipynb diff --git a/week2/nmap2.ipynb b/week2/nmap2.ipynb new file mode 100644 index 0000000..617b829 --- /dev/null +++ b/week2/nmap2.ipynb @@ -0,0 +1,889 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "guided-sense", + "metadata": {}, + "source": [ + "<h1 style=\"text-align: center; vertical-align: middle;\">Numerical Methods in Accelerator Physics</h1>\n", + "<h2 style=\"text-align: center; vertical-align: middle;\">Python examples -- Week 2</h2>" + ] + }, + { + "cell_type": "markdown", + "id": "instrumental-stylus", + "metadata": {}, + "source": [ + "<h2>Run this first!</h2>\n", + "\n", + "Imports and modules:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "universal-freedom", + "metadata": {}, + "outputs": [], + "source": [ + "from config2 import *\n", + "%matplotlib inline" + ] + }, + { + "cell_type": "markdown", + "id": "conventional-header", + "metadata": {}, + "source": [ + "<h3>Pendulum parameters, Hamiltonian</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cellular-passport", + "metadata": {}, + "outputs": [], + "source": [ + "m = 1 # point mass\n", + "g = 1 # magnitude of the gravitational field\n", + "L = 1 # length of the rod" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "worthy-bridges", + "metadata": {}, + "outputs": [], + "source": [ + "### Values of hamiltonian for comparison with theory\n", + "TH, PP = np.meshgrid(np.linspace(-np.pi * 1.1, np.pi * 1.1, 100), \n", + " np.linspace(-3, 3, 100))\n", + "\n", + "HH = hamiltonian(TH, PP)" + ] + }, + { + "cell_type": "markdown", + "id": "reverse-algebra", + "metadata": {}, + "source": [ + "<h3>Euler-Cromer Method (slide 10)</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "broke-edward", + "metadata": {}, + "outputs": [], + "source": [ + "def solve_eulercromer(theta, p, dt=0.1):\n", + " ### from Euler:\n", + " theta_next = theta + dt * p / (m * L*L)\n", + " p_next = p - dt * m * g * L * np.sin(theta_next)\n", + " return (theta_next, p_next)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "pediatric-prompt", + "metadata": {}, + "outputs": [], + "source": [ + "### Initial values:\n", + "theta_ini = -1.1\n", + "p_ini = 0\n", + "n_steps = 100" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "italic-muscle", + "metadata": {}, + "outputs": [], + "source": [ + "results_eulercromer = np.zeros((n_steps, 2), dtype=np.float32)\n", + "results_eulercromer[0] = (theta_ini, p_ini)\n", + "\n", + "for k in range(1, n_steps):\n", + " results_eulercromer[k] = solve_eulercromer(*results_eulercromer[k - 1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "leading-liberal", + "metadata": {}, + "outputs": [], + "source": [ + "### Explicit Euler for comparison:\n", + "results_euler = np.zeros((n_steps, 2), dtype=np.float32) # for comparison\n", + "results_euler[0] = (theta_ini, p_ini)\n", + "\n", + "for k in range(1, n_steps):\n", + " results_euler[k] = solve_euler(*results_euler[k - 1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "posted-irrigation", + "metadata": {}, + "outputs": [], + "source": [ + "plt.contour(TH, PP, HH, levels=10, linestyles='--', linewidths=1)\n", + "\n", + "plt.plot(results_euler[:, 0], results_euler[:, 1], c='b', label='Euler')\n", + "plt.plot(results_eulercromer[:, 0], results_eulercromer[:, 1], c='c', label='Euler-Cromer')\n", + "\n", + "plt.legend(bbox_to_anchor=(1.05, 1))\n", + "set_axes()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "scientific-comparison", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(\n", + " hamiltonian(results_euler[:, 0], results_euler[:, 1]), \n", + " c='b', label='Euler')\n", + "plt.plot(\n", + " hamiltonian(results_eulercromer[:, 0], results_eulercromer[:, 1]), \n", + " c='c', label='Euler-Cromer')\n", + "plt.axhline(hamiltonian(theta_ini, p_ini), c='r', label='theory')\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel(r'$\\mathcal{H}(\\theta, p)$')\n", + "plt.legend(bbox_to_anchor=(1.05, 1));" + ] + }, + { + "cell_type": "markdown", + "id": "fitted-occasions", + "metadata": {}, + "source": [ + "<h3>Leapfrog Method (slide 12)</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "lovely-silver", + "metadata": {}, + "outputs": [], + "source": [ + "def solve_leapfrog(theta, p, dt=0.1):\n", + " theta_half = theta + dt / 2 * p / (m * L*L)\n", + " p_next = p - dt * m * g * L * np.sin(theta_half)\n", + " theta_next = theta_half + dt / 2 * p_next / (m * L*L)\n", + " return (theta_next, p_next)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "civil-tunnel", + "metadata": {}, + "outputs": [], + "source": [ + "results_leapfrog = np.zeros((n_steps, 2), dtype=np.float32)\n", + "results_leapfrog[0] = (theta_ini, p_ini)\n", + "\n", + "for k in range(1, n_steps):\n", + " results_leapfrog[k] = solve_leapfrog(*results_leapfrog[k - 1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "determined-consortium", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(\n", + " hamiltonian(results_euler[:, 0], results_euler[:, 1]), \n", + " c='b', label='Euler, $\\mathcal{O}(\\Delta t^2)$')\n", + "plt.plot(\n", + " hamiltonian(results_eulercromer[:, 0], results_eulercromer[:, 1]), \n", + " c='c', label='Euler-Cromer, $\\mathcal{O}(\\Delta t^2)$')\n", + "plt.plot(\n", + " hamiltonian(results_leapfrog[:, 0], results_leapfrog[:, 1]), \n", + " c='k', label='leapfrog, $\\mathcal{O}(\\Delta t^3)$')\n", + "plt.axhline(hamiltonian(theta_ini, p_ini), c='r', lw=2, label='theory')\n", + "\n", + "plt.ylim(0.5, 0.6)\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel(r'$\\mathcal{H}(\\theta, p)$')\n", + "plt.legend(bbox_to_anchor=(1.05, 1));" + ] + }, + { + "cell_type": "markdown", + "id": "worst-bikini", + "metadata": {}, + "source": [ + "<h3>Collection of pendulums (slides 14 & 17)</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "respiratory-chaos", + "metadata": {}, + "outputs": [], + "source": [ + "theta_grid = np.linspace(-0.5 * np.pi, 0.5 * np.pi, 21)\n", + "p_grid = np.linspace(-0.3, 0.3, 5)\n", + "\n", + "thetas, ps = np.meshgrid(theta_grid, p_grid)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "smart-school", + "metadata": {}, + "outputs": [], + "source": [ + "N = len(theta_grid) * len(p_grid)\n", + "N" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "australian-baker", + "metadata": {}, + "outputs": [], + "source": [ + "plt.scatter(thetas, ps, c='b', marker='.')\n", + "\n", + "plot_hamiltonian()" + ] + }, + { + "cell_type": "markdown", + "id": "electrical-union", + "metadata": {}, + "source": [ + "<h3>Time evolution</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sixth-plain", + "metadata": {}, + "outputs": [], + "source": [ + "n_steps = 100" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "weird-lawyer", + "metadata": {}, + "outputs": [], + "source": [ + "results_thetas = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_thetas[0] = thetas.flatten()\n", + "\n", + "results_ps = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_ps[0] = ps.flatten()\n", + "\n", + "for k in range(1, n_steps):\n", + " results_thetas[k], results_ps[k] = solve_leapfrog(results_thetas[k - 1], results_ps[k - 1])" + ] + }, + { + "cell_type": "markdown", + "id": "special-canadian", + "metadata": {}, + "source": [ + "<h3>Observations of centroids (slide 19)</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "indonesian-redhead", + "metadata": {}, + "outputs": [], + "source": [ + "centroids_theta = 1/N * np.sum(results_thetas, axis=1)\n", + "centroids_p = 1/N * np.sum(results_ps, axis=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "outside-wonder", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(centroids_theta, label=r'$\\langle\\theta\\rangle$')\n", + "plt.plot(centroids_p, label=r'$\\langle p\\rangle$')\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('Centroid amplitude')\n", + "plt.legend(bbox_to_anchor=(1.05, 1));" + ] + }, + { + "cell_type": "markdown", + "id": "forbidden-merit", + "metadata": {}, + "source": [ + "<h3>Observations of RMS size</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "continental-preliminary", + "metadata": {}, + "outputs": [], + "source": [ + "var_theta = 1/N * np.sum(results_thetas * results_thetas, axis=1)\n", + "var_p = 1/N * np.sum(results_ps * results_ps, axis=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "metropolitan-princess", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(var_theta, label=r'$\\langle\\theta^2\\rangle$')\n", + "plt.plot(var_p, label=r'$\\langle p^2\\rangle$')\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('Variance')\n", + "plt.legend(bbox_to_anchor=(1.05, 1));" + ] + }, + { + "cell_type": "markdown", + "id": "quick-humor", + "metadata": {}, + "source": [ + "<h3>Distribution evolution over time</h3> " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "homeless-deployment", + "metadata": {}, + "outputs": [], + "source": [ + "k = 18\n", + "\n", + "plt.scatter(results_thetas[k], results_ps[k], c='b', marker='.')\n", + "\n", + "plot_hamiltonian()" + ] + }, + { + "cell_type": "markdown", + "id": "lined-glory", + "metadata": {}, + "source": [ + "<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", + "execution_count": null, + "id": "welsh-effort", + "metadata": {}, + "outputs": [], + "source": [ + "def emittance(theta, p):\n", + " N = len(theta)\n", + " \n", + " # subtract centroids\n", + " theta = theta - 1/N * np.sum(theta)\n", + " p = p - 1/N * np.sum(p)\n", + " \n", + " # compute Σ matrix entries\n", + " theta_sq = 1/N * np.sum(theta * theta)\n", + " p_sq = 1/N * np.sum(p * p)\n", + " crossterm = 1/N * np.sum(theta * p)\n", + " \n", + " # determinant of Σ matrix\n", + " epsilon = np.sqrt(theta_sq * p_sq - crossterm * crossterm)\n", + " return epsilon" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "impaired-russian", + "metadata": {}, + "outputs": [], + "source": [ + "results_emit = np.zeros(n_steps, dtype=np.float32)\n", + "\n", + "for k in range(n_steps):\n", + " results_emit[k] = emittance(results_thetas[k], results_ps[k])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "enhanced-dairy", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(results_emit)\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('RMS emittance $\\epsilon$');" + ] + }, + { + "cell_type": "markdown", + "id": "taken-survival", + "metadata": {}, + "source": [ + "<h3>Exercise: Linearized pendulum motion</h3>" + ] + }, + { + "cell_type": "markdown", + "id": "thorough-failing", + "metadata": {}, + "source": [ + "$$\\left.\\begin{matrix}\\,\n", + " \\theta_{k+1/2} &=& \\theta_k + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k}}{mL^2} \\\\\n", + " p_{k+1} &=& p_k - \\Delta t\\cdot m g L\\sin\\theta_{k+1/2} \\\\\n", + " \\theta_{k+1} &=& \\theta_{k+1/2} + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k+1}}{mL^2}\n", + "\\end{matrix}\\right\\} \\quad \\stackrel{\\text{Taylor}}{\\Longrightarrow} \\quad \\left\\{\\begin{matrix}\\,\n", + " \\theta_{k+1/2} &=& \\theta_k + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k}}{mL^2} \\\\\n", + " 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)} \\\\\n", + " \\theta_{k+1} &=& \\theta_{k+1/2} + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k+1}}{mL^2}\n", + "\\end{matrix}\\right.$$ " + ] + }, + { + "cell_type": "markdown", + "id": "quiet-updating", + "metadata": {}, + "source": [ + "$$\\left.\\begin{matrix}\\,\n", + " \\theta_{k+1/2} &=& \\theta_k + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k}}{mL^2} \\\\\n", + " p_{k+1} &=& p_k - \\Delta t\\cdot m g L\\sin\\theta_{k+1/2} \\\\\n", + " \\theta_{k+1} &=& \\theta_{k+1/2} + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k+1}}{mL^2}\n", + "\\end{matrix}\\right\\} \\quad \\stackrel{\\text{First order}}{\\Longrightarrow} \\quad \\left\\{\\begin{matrix}\\,\n", + " \\theta_{k+1/2} &=& \\theta_k + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k}}{mL^2} \\\\\n", + " p_{k+1} &=& p_k - \\Delta t\\cdot m g L\\cdot\\color{red}{\\theta_{k+1/2}} \\\\\n", + " \\theta_{k+1} &=& \\theta_{k+1/2} + {\\cfrac{\\Delta t}{2}}\\cdot\\cfrac{p_{k+1}}{mL^2}\n", + "\\end{matrix}\\right.$$ " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "living-spyware", + "metadata": {}, + "outputs": [], + "source": [ + "def solve_leapfrog_linear(theta, p, dt=dt):\n", + " theta_half = theta + dt / 2 * p / (m * L*L)\n", + " p_next = p - dt * m * g * L * theta_half\n", + " theta_next = theta_half + dt / 2 * p_next / (m * L*L)\n", + " return (theta_next, p_next)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bored-commodity", + "metadata": {}, + "outputs": [], + "source": [ + "results_thetas_linear = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_thetas_linear[0] = thetas.flatten()\n", + "\n", + "results_ps_linear = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_ps_linear[0] = ps.flatten()\n", + "\n", + "for k in range(1, n_steps):\n", + " results_thetas_linear[k], results_ps_linear[k] = solve_leapfrog_linear(results_thetas_linear[k - 1], results_ps_linear[k - 1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "quantitative-programmer", + "metadata": {}, + "outputs": [], + "source": [ + "var_theta_linear = 1/N * np.sum(results_thetas_linear * results_thetas_linear, axis=1)\n", + "var_p_linear = 1/N * np.sum(results_ps_linear * results_ps_linear, axis=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "three-better", + "metadata": { + "scrolled": true + }, + "outputs": [], + "source": [ + "plt.plot(var_theta_linear, label=r'$\\langle\\theta^2\\rangle$')\n", + "plt.plot(var_p_linear, label=r'$\\langle p^2\\rangle$')\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('Variance')\n", + "plt.legend(bbox_to_anchor=(1.05, 1));" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "statutory-thriller", + "metadata": {}, + "outputs": [], + "source": [ + "results_emit_linear = np.zeros(n_steps, dtype=np.float32)\n", + "\n", + "for k in range(n_steps):\n", + " results_emit_linear[k] = emittance(results_thetas_linear[k], results_ps_linear[k])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "equipped-campus", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(results_emit_linear)\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('RMS emittance $\\epsilon$');" + ] + }, + { + "cell_type": "markdown", + "id": "breeding-invalid", + "metadata": {}, + "source": [ + "<h3>Centroid offset in (original) pendulum</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "chemical-aquarium", + "metadata": {}, + "outputs": [], + "source": [ + "results_thetas2 = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_thetas2[0] = thetas.flatten() + 0.1 * np.pi\n", + "\n", + "results_ps2 = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_ps2[0] = ps.flatten()\n", + "\n", + "for k in range(1, n_steps):\n", + " results_thetas2[k], results_ps2[k] = solve_leapfrog(results_thetas2[k - 1], results_ps2[k - 1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "greek-brunswick", + "metadata": {}, + "outputs": [], + "source": [ + "centroids_theta2 = 1/N * np.sum(results_thetas2, axis=1)\n", + "centroids_p2 = 1/N * np.sum(results_ps2, axis=1)\n", + "\n", + "plt.plot(centroids_theta, label=r'no offset: $\\langle\\theta\\rangle$')\n", + "plt.plot(centroids_theta2, label=r'with offset: $\\langle\\theta\\rangle$')\n", + "plt.plot(centroids_p2, label=r'with offset: $\\langle p\\rangle$')\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('Centroid amplitude')\n", + "plt.legend(bbox_to_anchor=(1.05, 1));" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "voluntary-transparency", + "metadata": {}, + "outputs": [], + "source": [ + "results_emit2 = np.zeros(n_steps, dtype=np.float32)\n", + "\n", + "for k in range(n_steps):\n", + " results_emit2[k] = emittance(results_thetas2[k], results_ps2[k])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "legitimate-objective", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(results_emit, label='no offset')\n", + "plt.plot(results_emit2, label='with offset')\n", + "\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel('RMS emittance $\\epsilon$')\n", + "plt.legend();" + ] + }, + { + "cell_type": "markdown", + "id": "rough-funds", + "metadata": {}, + "source": [ + "<h3>Long-term behaviour</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "canadian-section", + "metadata": {}, + "outputs": [], + "source": [ + "n_steps = 3000" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "primary-lightning", + "metadata": {}, + "outputs": [], + "source": [ + "results_thetas3 = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_thetas3[0] = thetas.flatten() + 0.1 * np.pi\n", + "\n", + "results_ps3 = np.zeros((n_steps, N), dtype=np.float32)\n", + "results_ps3[0] = ps.flatten()\n", + "\n", + "for k in range(1, n_steps):\n", + " results_thetas3[k], results_ps3[k] = solve_leapfrog(results_thetas3[k - 1], results_ps3[k - 1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ordinary-advertising", + "metadata": {}, + "outputs": [], + "source": [ + "plot_macro_evolution(results_thetas3, results_ps3);" + ] + }, + { + "cell_type": "markdown", + "id": "figured-attendance", + "metadata": {}, + "source": [ + "<h3>Employ FFT on pendulum motion (slide 31)</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "coordinated-boring", + "metadata": {}, + "outputs": [], + "source": [ + "## Observe two particles:\n", + "i1 = N // 2 - 10\n", + "i2 = N // 2\n", + "\n", + "plt.plot(results_thetas3[:, i1])\n", + "plt.plot(results_thetas3[:, i2])\n", + "plt.xlabel('Steps $k$')\n", + "plt.ylabel(r'$\\theta$');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "supported-involvement", + "metadata": {}, + "outputs": [], + "source": [ + "## Obtain the FFT spectra for both:\n", + "spec1 = np.fft.rfft(results_thetas3[:, i1])\n", + "spec2 = np.fft.rfft(results_thetas3[:, i2])\n", + "\n", + "freq = np.fft.rfftfreq(n_steps)\n", + "\n", + "## Frequency of linearized system\n", + "freq_theory = np.sqrt(1 / 1) * 0.1 / (2 * np.pi)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "weird-retrieval", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(freq, np.abs(spec1))\n", + "plt.plot(freq, np.abs(spec2))\n", + "plt.axvline(freq_theory, c='k', zorder=0, ls='--')\n", + "plt.xlim(0, 0.05)\n", + "plt.xlabel('Phase advance per $\\Delta t$ [$2\\pi$]')\n", + "plt.ylabel('FFT spectrum')\n", + "plt.yscale('log');" + ] + }, + { + "cell_type": "markdown", + "id": "desirable-venice", + "metadata": {}, + "source": [ + "<h3>Frequency vs Amplitude</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "textile-commons", + "metadata": {}, + "outputs": [], + "source": [ + "specs = np.abs(np.fft.rfft(results_thetas3.T))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "standing-there", + "metadata": {}, + "outputs": [], + "source": [ + "max_ids = np.argmax(specs, axis=1)\n", + "amplitudes = np.max(specs, axis=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "developed-input", + "metadata": {}, + "outputs": [], + "source": [ + "i_range = np.where(ps.flatten() == 0)[0]\n", + "\n", + "plt.scatter(results_thetas3[0][i_range], freq[max_ids][i_range])\n", + "\n", + "plt.axhline(freq_theory, c='k', ls='--', zorder=0)\n", + "plt.xticks([-np.pi/2, 0, np.pi/2, np.pi], [r\"$-\\pi/2$\", \"0\", r\"$\\pi/2$\", r\"$\\pi$\"])\n", + "plt.xlabel(r'Initial $\\theta$ ($p=0$)')\n", + "plt.ylabel('Phase advance / \\n' + r'integration step $\\Delta t$ [$2\\pi$]');\n", + "plt.ylim(0.012, 0.0165);" + ] + }, + { + "cell_type": "markdown", + "id": "unavailable-scenario", + "metadata": {}, + "source": [ + "<h3>NAFF Algorithm (slide 32)</h3>" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "present-helena", + "metadata": {}, + "outputs": [], + "source": [ + "import PyNAFF" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "global-distributor", + "metadata": {}, + "outputs": [], + "source": [ + "freqs_naff = []\n", + "\n", + "for signal in results_thetas3.T[i_range]:\n", + " freq_naff = PyNAFF.naff(signal, turns=n_steps - 1, nterms=1)\n", + " try:\n", + " freq_naff = freq_naff[0, 1]\n", + " except IndexError:\n", + " freq_naff = 0\n", + " freqs_naff += [freq_naff]\n", + "\n", + "freqs_naff = np.array(freqs_naff)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "statutory-union", + "metadata": {}, + "outputs": [], + "source": [ + "plt.scatter(results_thetas3[0][i_range], freq[max_ids][i_range])\n", + "plt.plot(results_thetas3[0][i_range], freqs_naff, c='r', marker='.')\n", + "\n", + "plt.axhline(freq_theory, c='k', ls='--', zorder=0)\n", + "plt.xticks([-np.pi/2, 0, np.pi/2, np.pi], [r\"$-\\pi/2$\", \"0\", r\"$\\pi/2$\", r\"$\\pi$\"])\n", + "plt.xlabel(r'Initial $\\theta$ ($p=0$)')\n", + "plt.ylabel('Phase advance / \\n' + r'integration step $\\Delta t$ [$2\\pi$]');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "inappropriate-listening", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(results_thetas3[0][i_range], 100 * (freq[max_ids][i_range] - freqs_naff) / freqs_naff)\n", + "plt.xlabel(r'Initial $\\theta$ ($p=0$)')\n", + "plt.ylabel('Error (FFT - NAFF) / NAFF ' + r'[%]');" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.18" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} -- GitLab