{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import biot_savart_v4_3 as bs\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the size larger\n",
    "plt.rcParams[\"figure.figsize\"] = (10, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COIL_NAME = \"coil_6_spiral\"\n",
    "# COIL_NAME = \"coil_6_custom\"\n",
    "# COIL_NAME = \"coil_12_spiral\"\n",
    "COIL_NAME = \"coil_12_custom\"\n",
    "COIL = \"coils/\" + COIL_NAME + \".csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs.plot_coil(COIL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs.write_target_volume(\n",
    "    COIL,\n",
    "    COIL_NAME + \"volume.vol\",\n",
    "    (4, 4, 4),\n",
    "    (-2, -2, -2),\n",
    "    0.1,\n",
    "    0.1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in the results\n",
    "total_volume = bs.read_target_volume(COIL_NAME + \"volume.vol\")\n",
    "\n",
    "total_volume = -total_volume\n",
    "\n",
    "# reads the volume we created\n",
    "bs.plot_fields(total_volume, (4, 4, 4), (-2, -2, -2), 0.1, which_plane=\"z\", level=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "\n",
    "X, Y, Z = np.mgrid[-2:2:41j, -2:2:41j, -2:2:41j]\n",
    "# values = np.sqrt(total_volume[:,:,:,0]*total_volume[:,:,:,0] +\n",
    "# total_volume[:,:,:,1]*total_volume[:,:,:,1] +\n",
    "# total_volume[:,:,:,2]*total_volume[:,:,:,2])\n",
    "values = total_volume[:, :, :, 2]\n",
    "fig = go.Figure(\n",
    "    data=go.Volume(\n",
    "        x=X.flatten(),\n",
    "        y=Y.flatten(),\n",
    "        z=Z.flatten(),\n",
    "        value=np.abs(values.flatten()),\n",
    "        isomin=0.1,\n",
    "        isomax=0.5,\n",
    "        opacity=0.1,  # needs to be small to see through all surfaces\n",
    "        surface_count=17,  # needs to be a large number for good volume rendering\n",
    "    )\n",
    ")\n",
    "\n",
    "coil = bs.parse_coil(\"coils/\" + COIL_NAME + \".csv\")\n",
    "coil = coil[0:3, :]\n",
    "# add the coil as a line plot\n",
    "fig.add_trace(\n",
    "    go.Scatter3d(\n",
    "        x=coil[0, :],\n",
    "        y=coil[1, :],\n",
    "        z=coil[2, :],\n",
    "        mode=\"lines\",\n",
    "        line=dict(color=\"red\", width=5),\n",
    "    )\n",
    ")\n",
    "# make the plotly graph bigger\n",
    "fig.update_layout(\n",
    "    autosize=False,\n",
    "    width=1400,\n",
    "    height=1000,\n",
    ")\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice = [total_volume[:, 21, :, i].T for i in range(3)]\n",
    "x_contour = slice[0]\n",
    "y_contour = slice[1]\n",
    "z_contour = slice[2]\n",
    "\n",
    "print(x_contour.shape)\n",
    "print(y_contour.shape)\n",
    "print(z_contour.shape)\n",
    "\n",
    "x = np.arange(0, x_contour.shape[0], 1)\n",
    "y = np.arange(0, z_contour.shape[1], 1)\n",
    "xG, yG = np.meshgrid(x, y)\n",
    "\n",
    "# calclulate the magnitude of the field\n",
    "mag = np.sqrt(x_contour ** 2 + y_contour ** 2 + z_contour ** 2)\n",
    "\n",
    "# color = 2 * np.log(np.hypot(x_contour, z_contour))\n",
    "color = np.log(mag)\n",
    "stream_plot = plt.streamplot(\n",
    "    x,\n",
    "    y,\n",
    "    x_contour,\n",
    "    z_contour,\n",
    "    color=color,\n",
    "    linewidth=1,\n",
    "    cmap=plt.cm.Reds,\n",
    "    density=2,\n",
    "    arrowstyle=\"->\",\n",
    "    arrowsize=1.5,\n",
    ")\n",
    "# set the axis to be equal\n",
    "plt.axis(\"equal\")\n",
    "\n",
    "# set the plt size to 5x5\n",
    "plt.rcParams[\"figure.figsize\"] = (10, 10)\n",
    "\n",
    "# save the figure\n",
    "plt.savefig(COIL_NAME + \"_field_lines.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.cm as cm\n",
    "\n",
    "# get the total left and right field in the X direction\n",
    "slice = [total_volume[:, :, 21, i].T for i in range(3)]\n",
    "x_contour = slice[0]\n",
    "y_contour = slice[1]\n",
    "z_contour = slice[2]\n",
    "\n",
    "# vmin = np.min([np.min(x_contour), np.min(z_contour), np.min(y_contour)])\n",
    "# vmax = np.max([np.max(x_contour), np.max(z_contour), np.max(y_contour)])\n",
    "\n",
    "vmin = np.min(x_contour)\n",
    "vmax = np.max(x_contour)\n",
    "\n",
    "# -7.970943527635176 7.775784181783158\n",
    "vmin = -7.970943527635176\n",
    "vmax = 7.775784181783158\n",
    "\n",
    "print(vmin, vmax)\n",
    "\n",
    "left_x = np.sum(x_contour[x_contour < 0])\n",
    "right_x = np.sum(x_contour[x_contour > 0])\n",
    "\n",
    "print(\"Left X: \", left_x)\n",
    "print(\"Right X: \", right_x)\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (5, 5)\n",
    "\n",
    "# make the plot larger\n",
    "plt.rcParams[\"figure.figsize\"] = (10, 10)\n",
    "\n",
    "# plot the coil\n",
    "coil = bs.parse_coil(\"coils/\" + COIL_NAME + \".csv\")\n",
    "coil = coil[0:3, :]\n",
    "coil = coil * 10 + 20\n",
    "plt.plot(coil[0, :], coil[1, :], color=\"black\", alpha=0.1)\n",
    "\n",
    "# plt.imshow(x_contour, cmap='hot', interpolation='nearest')\n",
    "plt.contourf(\n",
    "    y_contour,\n",
    "    vmin=vmin,\n",
    "    vmax=vmax,\n",
    "    cmap=cm.magma,\n",
    "    levels=50,\n",
    ")\n",
    "\n",
    "# hide the axis\n",
    "plt.axis(\"off\")\n",
    "\n",
    "# remove any padding\n",
    "plt.tight_layout()\n",
    "plt.margins(x=0)\n",
    "plt.margins(x=0)\n",
    "\n",
    "plt.savefig(COIL_NAME + \"_field_y.png\", bbox_inches=\"tight\", pad_inches=0)\n",
    "\n",
    "print(vmin, vmax)\n",
    "\n",
    "# 6 coil custom\n",
    "# Left X:  -565.6790142743889\n",
    "# Right X:  544.1011578804321\n",
    "\n",
    "# 6 coil spiral\n",
    "# Left X:  -512.1107132193335\n",
    "# Right X:  534.1946450523811\n",
    "\n",
    "# 12 coil spiral\n",
    "# Left X:  -218.54696343677543\n",
    "# Right X:  233.84835661493196\n",
    "\n",
    "# 12 coil custom\n",
    "# Left X:  -361.95894587646944\n",
    "# Right X:  347.13174239269904"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the magitude of the slice\n",
    "plt.imshow(np.log(mag), cmap=\"hot\", interpolation=\"nearest\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice = [total_volume[:, :, 25, i].T for i in range(3)]\n",
    "x_contour = slice[0]\n",
    "y_contour = slice[1]\n",
    "z_contour = slice[2]\n",
    "\n",
    "mag = np.sqrt(x_contour ** 2 + y_contour ** 2 + z_contour ** 2)\n",
    "\n",
    "print(np.max(mag), np.min(mag))\n",
    "\n",
    "# plot a histogram of the magnitude\n",
    "# plt.hist(mag.flatten(), bins=10)\n",
    "\n",
    "plt.hist(z_contour.flatten(), bins=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the z contour as a 3d surface plot\n",
    "fig = plt.figure()\n",
    "\n",
    "# change the figure size to 16.9\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "\n",
    "xG_ = xG[10:30, 10:30]\n",
    "yG_ = yG[10:30, 10:30]\n",
    "z_contour_ = z_contour[10:30, 10:30]\n",
    "\n",
    "print(xG.shape, yG.shape, z_contour.shape)\n",
    "\n",
    "# change the axis limits to 10 to 30\n",
    "ax.set_xlim(10, 30)\n",
    "ax.set_ylim(10, 30)\n",
    "surf = ax.plot_surface(\n",
    "    xG_, yG_, z_contour_, cmap=cm.plasma, linewidth=0, antialiased=False, alpha=0.5\n",
    ")\n",
    "# add the coil to the plot as a set of lines - the X and Y coordinates of the coil are in the first two elements of the coil array\n",
    "coil = bs.parse_coil(\"coils/\" + COIL_NAME + \".csv\")\n",
    "coil = coil[0:2, :]\n",
    "# need to scale the coil by 10 and shift it by 20, 20\n",
    "coil = coil * 10 + 20\n",
    "# draw the coil in yellow\n",
    "ax.plot(coil[0, :], coil[1, :], 0, color=\"black\")\n",
    "\n",
    "# hide the axes\n",
    "ax.set_axis_off()\n",
    "\n",
    "# save to png\n",
    "plt.savefig(COIL_NAME + \"_field_surface_z.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write out the coil as SVG\n",
    "coil = bs.parse_coil(\"coils/\" + COIL_NAME + \".csv\")\n",
    "coil = coil[0:2, :]\n",
    "with open(COIL_NAME + \"_coil.svg\", \"w\") as f:\n",
    "    # write the header\n",
    "    f.write('<svg viewBox=\"0 0 400 400\" xmlns=\"http://www.w3.org/2000/svg\">')\n",
    "    # the X and Y points of the coil the X and Y coordinates of the coil are in the first two elements of the coil array\n",
    "    # write out the coil as a set of lines\n",
    "    print(len(coil[0]) - 1)\n",
    "    # write out the coil as a polyline\n",
    "    f.write('<polyline points=\"')\n",
    "    for i in range(0, len(coil[0]), 10):\n",
    "        f.write(str(coil[0, i] * 100 + 200) + \",\" + str(coil[1, i] * 100 + 200) + \" \")\n",
    "    f.write('\" style=\"fill:none;stroke:black;stroke-width:1\" />')\n",
    "    # write the footer\n",
    "    f.write(\"</svg>\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.7 ('venv': venv)",
   "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.10.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "fc384f9db26c31784edfba3761ba3d2c7b2f9b8a63e03a9eb0778fc35334efe1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}