{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8d06107",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"\n",
    "Created on Tue Aug 24 16:24:52 2021\n",
    "\n",
    "@author: Joschka\n",
    "\"\"\"\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sys\n",
    "sys.path.insert(0,'..')\n",
    "\n",
    "from src import coil_class as BC\n",
    "from src import physical_constants as cs\n",
    "\n",
    "#from IPython import get_ipython\n",
    "#get_ipython().run_line_magic('matplotlib', 'qt')\n",
    "#get_ipython().run_line_magic('matplotlib', 'inline')\n",
    "\n",
    "\n",
    "from scipy.optimize import curve_fit\n",
    "plt.rcParams['axes.grid'] = True\n",
    "plt.rcParams['grid.alpha'] = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99eb6ef1-8920-4f7d-9b58-fd00c731d2bc",
   "metadata": {},
   "source": [
    "## Set up coils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "199e2602-bec6-429a-a639-6ea35ababb86",
   "metadata": {},
   "outputs": [],
   "source": [
    "HH = 1\n",
    "I = 5\n",
    "\n",
    "d_coils = 50\n",
    "Radius = 30\n",
    "\n",
    "layers = 4\n",
    "windings = 4\n",
    "wire_width = 1\n",
    "wire_height = 1\n",
    "\n",
    "Coil = BC.BCoil(HH,d_coils,Radius, layers, windings, wire_width, wire_height)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40cee79d-d6ef-4df4-ae03-5e82785433e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.arange(-100,100,10)\n",
    "print(np.size(z))\n",
    "r = np.arange(1e-15,100,5)\n",
    "\n",
    "B, B_tot = Coil.B_multiple_3d(I, r,z,raster = 4)\n",
    "\n",
    "z_m, r_m = np.meshgrid(z,r)\n",
    "\n",
    "#plt.figure(figsize=(16,10))\n",
    "plt.quiver(r_m,z_m,B[:,:,1],B[:,:,0])\n",
    "plt.xlabel(\"radius r [m]\")\n",
    "plt.ylabel(\"z-axis [m]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28dcdd21-f4b6-49e6-a935-4012a941f186",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.arange(-100,100,1)\n",
    "r = np.arange(1e-3,100,1)\n",
    "\n",
    "B, B_tot = Coil.B_multiple_3d(I, r,z,raster = 2)\n",
    "\n",
    "z_m, r_m = np.meshgrid(z,r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad02e8cd-e92a-4d27-845d-cc7e22d6e7e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.concatenate((-np.flip(r),r))\n",
    "B_tot_x = np.concatenate((np.flip(B_tot[:,len(z)//2]),B_tot[:,len(z)//2]))\n",
    "B_tot_z = B_tot[0,:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52475c2c-e2f0-4615-a5f4-9f0a86155a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x,B_tot_x)\n",
    "plt.xlabel(\"radius r [mm]\")\n",
    "plt.ylabel(\"total field B_tot [G]\")\n",
    "#plt.xlim(-0.1,0.1)\n",
    "plt.show()\n",
    "plt.plot(z,B_tot_z)\n",
    "plt.xlabel(\"z_axis [mm]\")\n",
    "plt.ylabel(\"total field B_tot [G]\")\n",
    "plt.show()\n",
    "\n",
    "\"\"\"\n",
    "plt.plot(x,np.gradient(B_tot_x,x))\n",
    "plt.xlabel(\"radius r [mm]\")\n",
    "plt.ylabel(\"total field B_tot [G]\")\n",
    "#plt.xlim(0,0.01)\n",
    "plt.show()\n",
    "plt.plot(z,np.gradient(B_tot[0,:],z))\n",
    "plt.xlabel(\"z_axis [mm]\")\n",
    "plt.ylabel(\"total field B_tot [G]\")\n",
    "plt.show()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37c71478-cae2-40ea-8712-b2f751c6c54d",
   "metadata": {},
   "source": [
    "## Fit harmonic function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "443f51a2-a502-4d4a-8692-15610c193cc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(x,a,b):\n",
    "    return a*x**2+b\n",
    "\n",
    "x_SI = 1e-3 * x\n",
    "z_SI = 1e-3 * z\n",
    "\n",
    "B_tot_x_SI = 1e-4*B_tot_x\n",
    "B_tot_z_SI = 1e-4*B_tot_z\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59145c9b-aac3-4bea-ba79-b2a3992d05f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nr_points = 20\n",
    "a = 100-nr_points//2\n",
    "b = 100+ nr_points//2\n",
    "popt_x, pcov = curve_fit(func,x_SI[a:b],B_tot_x_SI[a:b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c6fb78-ba74-41e3-9a2a-c4dca7109643",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(1e3*x_SI,1e4*func(x_SI,popt_x[0],popt_x[1]),label = \"harmonic fit\")\n",
    "plt.plot(1e3*x_SI,1e4*B_tot_x_SI, label = r\"$ B_{tot} = \\sqrt{r^2 + z^2} $\" )\n",
    "#plt.ylim(0,8)\n",
    "#plt.xlim(-0.01,0.01)\n",
    "plt.xlabel(\"x-axis [mm]\")\n",
    "plt.xlim(-50,50)\n",
    "plt.ylim(0,20)\n",
    "plt.ylabel(\"B_tot [G]\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c981bb-02be-4de9-8990-849b6a7fbee1",
   "metadata": {},
   "outputs": [],
   "source": [
    "nr_points = 20\n",
    "a = 100-nr_points//2\n",
    "b = 100+ nr_points//2\n",
    "popt_z, pcov = curve_fit(func,z_SI[a:b],B_tot_z_SI[a:b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f7cdd7-f13c-41e2-be5e-a375dc5e4d06",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(1e3*z_SI,1e4*func(z_SI,popt_z[0],popt_z[1]),label = \"harmonic fit\")\n",
    "plt.plot(1e3*z_SI,1e4*B_tot_z_SI, label = r\"$ B_{tot} = \\sqrt{r^2 + z^2} $\")\n",
    "plt.xlabel(\"x-axis [mm]\")\n",
    "plt.xlim(-50,50)\n",
    "plt.ylim(0,50)\n",
    "plt.ylabel(\"B_tot [G]\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "624fbda5-4658-4f9e-ba0f-33168c27e5c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 2.69e-25\n",
    "\n",
    "k_x = -2*popt_x[0]*9.9*cs.mu_B\n",
    "\n",
    "omega = np.sqrt(k_x/m)\n",
    "f = omega/(2*np.pi)\n",
    "\n",
    "T = 1/f\n",
    "T_exp = T/4\n",
    "print(f\"T_expansion = T/4 = {T_exp*1e3} ms\" )\n",
    "print(f\"omega_x = {omega}\") \n",
    "\n",
    "t_tof = 20e-3\n",
    "M = omega * t_tof\n",
    "\n",
    "print(f\"Magnification for t_tof = {t_tof*1e3} ms: M = {M}\")\n",
    "\n",
    "\n",
    "start_z = 1e-6\n",
    "d_t = 1e-3\n",
    "def force(z):\n",
    "    return 2*0.248*z*9.9*cs.mu_B\n",
    "z = start_z\n",
    "v = 0\n",
    "for t in np.arange(0,T_exp,d_t):\n",
    "    v = v + force(z)/m * d_t\n",
    "    #print(v)\n",
    "    z = z + v * d_t\n",
    "print(f\"for z_start = 1 μm after T_expansion z_end = {z*1e6} μm\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1561200c-5630-4c8a-b0b5-b6a535805eef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}