{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ebf53fc1",
   "metadata": {},
   "source": [
    "# Basic Coils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001294fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax_sbgeom as jsb\n",
    "import jax.numpy as jnp\n",
    "import jax \n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "import pyvista\n",
    "import plotly.graph_objects as go"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96fade3b",
   "metadata": {},
   "source": [
    "We first load some proprietary data, but any file will work with the shape [n_coils, n_position, 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8aa631",
   "metadata": {},
   "outputs": [],
   "source": [
    "coil_positions = jnp.load(\"helias5_coils.npy\")\n",
    "print(coil_positions.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b36043cb",
   "metadata": {},
   "source": [
    "We can create a discrete coil from this and calculate some position and tangent from this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf2318f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sbgeom.coils import DiscreteCoil, CoilSet\n",
    "coil = DiscreteCoil.from_positions(coil_positions[0])\n",
    "print(\"Position\", coil.position(0.2))\n",
    "print(\"Tangent\", coil.tangent(0.2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b1a03e6",
   "metadata": {},
   "source": [
    "And plot the coil in 3D space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb748bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = jnp.linspace(0,1,100)\n",
    "coil_s = coil.position(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c4077c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure(layout=go.Layout(width=600, height=500))\n",
    "fig.add_trace(go.Scatter3d(x = coil_s[:,0], y= coil_s[:,1], z=coil_s[:,2], mode='lines', name=\"coil\", marker=dict(color='red', size=5)))\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cd2125f",
   "metadata": {},
   "source": [
    "A discrete coil has some limitations: it does not have a well-defined curvature and therefore normal for example. This can complicate creating a finite size somewhat, so we first convert to Fourier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8429c53",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Discrete coil normal: \" , coil.normal(0.2))\n",
    "fourier_coil = jsb.coils.convert_to_fourier_coil(coil)\n",
    "print(\"Fourier coil normal: \" , fourier_coil.normal(0.2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46a7a66d",
   "metadata": {},
   "source": [
    "We can create a finite size in some different ways, we'll plot a meshed version of them side-by-side.\n",
    "\n",
    "A finite size is defined by a definition of a radial vector. Together with the tangent, this creates an orthonormal frame (the last is obtained by a cross product!).\n",
    "\n",
    "Frenet-Serret uses the normal as radial vector, the centroid frame uses the normalised vector from the point to the centre, and RotationMinimizedFrame uses a frame that minimizes rotation.\n",
    "\n",
    "Extra arguments needed for a FiniteSizeCoil have to be appended as can be seen in the rotation-minimized coil (here it means the number of sample points for creating a rotation miminized frame)\n",
    "\n",
    "User defined radial vectors are also supported using RadialVectorFrame (of which RotationMinimizedFrame is in fact a subclass!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f6cc614",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sbgeom.coils import FrenetSerretFrame, CentroidFrame, RotationMinimizedFrame, FiniteSizeCoil\n",
    "\n",
    "frenet_coil   = FiniteSizeCoil.from_coil(fourier_coil, FrenetSerretFrame)\n",
    "rmf_coil      = FiniteSizeCoil.from_coil(fourier_coil, RotationMinimizedFrame, 100)\n",
    "centroid_coil = FiniteSizeCoil.from_coil(fourier_coil, CentroidFrame)\n",
    "\n",
    "mesh_frenet = jsb.coils.mesh_coil_surface(frenet_coil, width_radial = 0.3, width_phi = 0.3, n_s = 100)\n",
    "mesh_rmf = jsb.coils.mesh_coil_surface(rmf_coil, width_radial = 0.3, width_phi = 0.3, n_s = 100)\n",
    "mesh_centroid = jsb.coils.mesh_coil_surface(centroid_coil, width_radial = 0.3, width_phi = 0.3, n_s = 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0982b972",
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "from plotly.subplots import make_subplots\n",
    "import numpy as np\n",
    "def add_mesh_to_plotly(fig, mesh, mesh_kwargs, trace_kwargs = {}):\n",
    "    x,y,z = mesh[0][:,0], mesh[0][:,1], mesh[0][:,2]\n",
    "    i,j,k = mesh[1][:,0], mesh[1][:,1], mesh[1][:,2]\n",
    "    fig.add_trace(go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, showlegend=True, **mesh_kwargs), **trace_kwargs)\n",
    "\n",
    "\n",
    "# create 1x2 grid of 3D plots\n",
    "fig = make_subplots(\n",
    "    rows=1,\n",
    "    cols=3,\n",
    "    specs=[[{\"type\": \"scene\"}, {\"type\": \"scene\"}, {\"type\": \"scene\"}]],\n",
    "    subplot_titles=(\"Plot 1\", \"Plot 2\", \"Plot 3\"),\n",
    ")\n",
    "add_mesh_to_plotly(fig, mesh_frenet, mesh_kwargs={\"color\": \"white\", \"opacity\": 1.0, }, trace_kwargs={\"row\":1, \"col\":1})\n",
    "\n",
    "add_mesh_to_plotly(fig, mesh_rmf, mesh_kwargs={\"color\": \"lightblue\", \"opacity\": 1.0, \"name\": \"RMF\"}, trace_kwargs={\"row\":1, \"col\":2})\n",
    "add_mesh_to_plotly(fig, mesh_centroid, mesh_kwargs={\"color\": \"lightgreen\", \"opacity\": 1.0, \"name\": \"Centroid\"}, trace_kwargs={\"row\":1, \"col\":3})\n",
    "\n",
    "fig.update_layout(height=500, width=1200)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee715284",
   "metadata": {},
   "source": [
    "The Frenet coils is rather twisty, and is therefore not very useful in this case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0f9125d",
   "metadata": {},
   "source": [
    "All these functions also work with a coilset, so let's plot a full coilset with the RotationMinimizedFrame. Note we are now using the DiscreteCoil as a batched dataclass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44680e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sbgeom.coils import FiniteSizeCoilSet\n",
    "coilset = CoilSet(coils = DiscreteCoil.from_positions(coil_positions))\n",
    "fourier_coilset = jsb.coils.convert_to_fourier_coilset(coilset)\n",
    "\n",
    "rmf_coilset = FiniteSizeCoilSet.from_coilset(fourier_coilset, RotationMinimizedFrame, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd65ffca",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "mesh = jsb.coils.mesh_coilset_surface(rmf_coilset, width_radial = 0.3, width_phi = 0.3, n_s = 100)\n",
    "add_mesh_to_plotly(fig, mesh, mesh_kwargs={\"color\": \"lightblue\", \"opacity\": 1.0, \"name\": \"RMF\"}, trace_kwargs={})\n",
    "fig.update_layout(dict(scene=dict(aspectmode='data')), width = 1200, height=700)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08c6c251",
   "metadata": {},
   "source": [
    "### Coil winding surface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36449adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_wireframe_to_plotly(fig, mesh, **kwargs):\n",
    "    import numpy as onp\n",
    "    triangles = mesh[0][mesh[1]]\n",
    "    list_x = onp.full((triangles.shape[0], triangles.shape[1] + 1), np.nan)\n",
    "    list_y = onp.full((triangles.shape[0], triangles.shape[1] + 1), np.nan)\n",
    "    list_z = onp.full((triangles.shape[0], triangles.shape[1] + 1), np.nan)\n",
    "\n",
    "    list_x[:,:-1] = triangles[:, :, 0]\n",
    "    list_y[:,:-1] = triangles[:, :, 1]\n",
    "    list_z[:,:-1] = triangles[:, :, 2]\n",
    "\n",
    "    fig.add_trace(go.Scatter3d(x= list_x.ravel(), y = list_y.ravel(), z = list_z.ravel(), mode = 'lines', line =dict(color='black', width=3), showlegend=False), **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5715bd57",
   "metadata": {},
   "source": [
    "We  can generate a coil winding surface by connecting adjacent coils:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a7e71b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cws = jsb.coils.create_coil_winding_surface(fourier_coilset, n_points_per_coil= 50, n_points_phi = 100, surface_type=\"spline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4900a0c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "\n",
    "add_mesh_to_plotly(fig, mesh, mesh_kwargs={\"color\": \"lightblue\", \"opacity\": 1.0, \"name\": \"RMF\"}, trace_kwargs={})\n",
    "add_mesh_to_plotly(fig, cws, mesh_kwargs={\"color\": \"lightgreen\", \"opacity\": 1.0, \"name\": \"Coil Winding\"}, trace_kwargs={})\n",
    "add_wireframe_to_plotly(fig, cws)\n",
    "fig.update_layout(dict(scene=dict(aspectmode='data')), width = 1200, height=700)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38e341d4",
   "metadata": {},
   "source": [
    "We can see that the resulting surface looks rather wiggly, with some straight edges and self-intersections.  We can optimize it so it doesn't have that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae1c2c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "cws_opt = jsb.coils.create_optimized_coil_winding_surface(fourier_coilset, n_points_per_coil= 50, n_points_phi = 100, surface_type=\"spline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bb2467",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "\n",
    "add_mesh_to_plotly(fig, mesh, mesh_kwargs={\"color\": \"lightblue\", \"opacity\": 1.0, \"name\": \"RMF\"}, trace_kwargs={})\n",
    "add_mesh_to_plotly(fig, cws_opt, mesh_kwargs={\"color\": \"lightgreen\", \"opacity\": 1.0, \"name\": \"Coil Winding\"}, trace_kwargs={})\n",
    "add_wireframe_to_plotly(fig, cws_opt)\n",
    "fig.update_layout(dict(scene=dict(aspectmode='data')), width = 1200, height=700)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad5833f4",
   "metadata": {},
   "source": [
    "Much better!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa0b006",
   "metadata": {},
   "source": [
    "We can obtain a thickness matrix from the LCFS using the built-in BVH ray-tracing (this is by far not as fast as a proper ray-tracer, but it works fine)\n",
    "\n",
    "(The longest time is actually spent on compiling the BVH generation and raytracing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b73a2a69",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux_surface_no_phi = jsb.flux_surfaces.FluxSurfaceNormalExtendedNoPhi.from_hdf5(\"helias5_vmec.nc4\")\n",
    "theta_mg, phi_mg, thickness_matrix    = jsb.flux_surfaces.generate_thickness_matrix(flux_surface_no_phi, cws_opt, n_theta = 200, n_phi = 300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b429e760",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.pcolormesh(phi_mg, theta_mg, thickness_matrix)\n",
    "plt.xlabel(\"$\\\\phi$\")\n",
    "plt.ylabel(\"$\\\\theta$\")\n",
    "plt.colorbar(label=\"Thickness\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a62fc08",
   "metadata": {},
   "source": [
    "Using this thickness matrix, we can convert it to a VMEC representation, with equal arclength! This allows us to connect the mesh easily to a plasma mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d76a569",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux_surface_extended = jsb.flux_surfaces.create_extended_flux_surface_d_interp_equal_arclength(flux_surface_no_phi, thickness_matrix, n_theta = 50, n_phi = 60, n_theta_s_arclength=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "721c8e37",
   "metadata": {},
   "outputs": [],
   "source": [
    "extended_mesh = jsb.flux_surfaces.mesh_surfaces_closed(flux_surface_extended, 1.0, 2.0, jsb.flux_surfaces.ToroidalExtent.full_module(flux_surface_extended), 50, 60, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "226fa1d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "\n",
    "add_mesh_to_plotly(fig, mesh, mesh_kwargs={\"color\": \"lightblue\", \"opacity\": 1.0, \"name\": \"RMF\"})\n",
    "add_mesh_to_plotly(fig, extended_mesh, mesh_kwargs={\"color\": \"lightgreen\", \"opacity\": 1.0, \"name\": \"Extended Flux Surface\"})\n",
    "add_wireframe_to_plotly(fig, extended_mesh)\n",
    "fig.update_layout(height=500, width=1200)\n",
    "fig.update_layout(dict(scene=dict(aspectmode='data')), width = 1200, height=700)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "stellsim",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}