{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "101a5024",
   "metadata": {},
   "source": [
    "### Basic Flux Surfaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f58471c",
   "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": "e3c0b049",
   "metadata": {},
   "source": [
    "We first import the data from some VMEC file (we use a proprietary file here, but any can be used after converting to netcdf4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b10d2042",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sbgeom.flux_surfaces import FluxSurface\n",
    "flux_surface = FluxSurface.from_hdf5(\"helias5_vmec.nc4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbd0e73f",
   "metadata": {},
   "source": [
    "Then, we can obtain positions from it easily:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462fa52d",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux_surface.cartesian_position(s = 1.0,  theta = 0.2, phi = 0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f47cf736",
   "metadata": {},
   "source": [
    "This base class does not go beyond the LCFS (same output):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b6d0601",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux_surface.cartesian_position(s = 2.0,  theta = 0.2, phi = 0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efbdd818",
   "metadata": {},
   "source": [
    "We can plot a poloidal slice:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38da3dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_poloidal_slice(fig, flux_surface, phi, s_values, n_theta = 50, line_kwargs = {}):\n",
    "    theta = jnp.linspace(0, 2 * jnp.pi, n_theta)\n",
    "    s_mg, theta_mg = jnp.meshgrid(jnp.array(s_values), theta, indexing=\"ij\")\n",
    "    cylindrical_values = flux_surface.cylindrical_position(s_mg, theta_mg, phi)\n",
    "    \n",
    "    for i in range(cylindrical_values.shape[0]):\n",
    "        line_kwargs['name'] = f\"s = {s_values[i]:.2f}\"\n",
    "        fig.add_trace(go.Scatter(x=cylindrical_values[i, :,0], y=cylindrical_values[i, :,1], mode='lines', **line_kwargs))\n",
    "\n",
    "fig = go.Figure(layout=go.Layout(width=600, height=500))\n",
    "plot_poloidal_slice(fig, flux_surface, phi = 0.3, s_values = [0.2, 0.5, 0.8, 1.0, 2.0], line_kwargs = {\"line\": dict(width=2)})\n",
    "fig.update_xaxes(scaleanchor=\"y\", scaleratio=1)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3846a7d9",
   "metadata": {},
   "source": [
    "This doesn't have the capability yet to ge beyond the LCFS, as you can see in the figure above.\n",
    "\n",
    "To do this, we use a different class, that extends the geometry by using the normal vector of the LCFS and using $s-1$ as the distance to the LCFS.\n",
    "We can obtain this by a convenience method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd07a599",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sbgeom.flux_surfaces import FluxSurfaceNormalExtended\n",
    "\n",
    "flux_surface_normal_extended = FluxSurfaceNormalExtended.from_flux_surface(flux_surface)\n",
    "\n",
    "fig = go.Figure(layout=go.Layout(width=600, height=500))\n",
    "plot_poloidal_slice(fig,flux_surface_normal_extended, phi = 0.2, s_values = [0.1,0.5,0.9,1.0, 2.0])\n",
    "fig.update_xaxes(scaleanchor=\"y\", scaleratio=1)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef4b7bea",
   "metadata": {},
   "source": [
    "This in principle works, but a problem is that the normal vector has a $\\phi$ component. In VMEC parametrisation, normally, $\\phi_{in} = \\phi_{out}$, which is quite convenient.\n",
    "\n",
    "To show this, we use the mesh functionality to mesh a surface beyond the LCFS and show that it intersects some constant-$\\phi$ planes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a09d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_sbgeom.flux_surfaces import ToroidalExtent\n",
    "def Plot_Phi_Plane(phi, fig, rmin,rmax, zmin,zmax, alpha = 1, color = 'blue', **kwargs):\n",
    "    '''\n",
    "    Plots a simple plane\n",
    "    '''\n",
    "    x = [rmin * jnp.cos(phi), rmax * jnp.cos(phi), rmax * jnp.cos(phi), rmin * jnp.cos(phi), rmin * jnp.cos(phi)]\n",
    "    y = [rmin * jnp.sin(phi), rmax * jnp.sin(phi), rmax * jnp.sin(phi), rmin * jnp.sin(phi), rmin * jnp.sin(phi)]\n",
    "    z = [zmin, zmin, zmax, zmax, zmin]\n",
    "        \n",
    "    fig = fig.add_trace(go.Scatter3d(x=x, y=y, z=z, surfaceaxis=1, opacity=alpha, surfacecolor=color, mode='lines', line=dict(color='black', width=1), showlegend=False, name=\"phi plane\"), **kwargs)\n",
    "    fig = fig.add_trace(go.Scatter3d(x=x, y=y, z=z, mode='lines', line=dict(color='black', width=3), showlegend=False), **kwargs)\n",
    "\n",
    "def create_figure_planes(toroidal_extent):\n",
    "    fig = go.Figure()\n",
    "    fig.update_layout(autosize=False, width=600, height=500, margin=dict(l=0, r=0, b=0, t=0))\n",
    "    Plot_Phi_Plane(phi = toroidal_extent.start, fig = fig, rmin = 16, rmax=28, zmin=-6, zmax=6, alpha = 0.7,  color = 'lightblue')\n",
    "    Plot_Phi_Plane(phi = toroidal_extent.end  , fig = fig, rmin = 16, rmax=28, zmin=-6, zmax=6, alpha = 0.7,  color = 'lightblue')\n",
    "    return fig\n",
    "\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",
    "fig = create_figure_planes(ToroidalExtent.half_module(flux_surface_normal_extended))\n",
    "lcfs_mesh = jsb.flux_surfaces.mesh_surface(flux_surface_normal_extended, s = 1.0, \n",
    "                                           toroidal_extent=ToroidalExtent.full_module(flux_surface_normal_extended, -0.25 * 2 * jnp.pi / flux_surface_normal_extended.nfp),\n",
    "                                            n_theta = 50, \n",
    "                                            n_phi = 100)\n",
    "d_2_mesh = jsb.flux_surfaces.mesh_surface(flux_surface_normal_extended, s = 3.0, toroidal_extent=ToroidalExtent.half_module(flux_surface_normal_extended), n_theta = 50, n_phi = 100)\n",
    "add_mesh_to_plotly(fig, lcfs_mesh, dict(opacity=0.3, color='pink', name = \"LCFS\"))\n",
    "add_mesh_to_plotly(fig, d_2_mesh, dict(opacity=1.0, color='lightblue', name = \"d=2.0\"))\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa5eb01",
   "metadata": {},
   "source": [
    "Clearly, the surface intersects the half module planes.\n",
    "\n",
    "To remedy this, we can project the normal vector on the RZ plane and use that to parametrise distance beyond the LCFS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e80c8f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_same_phi(flux_surface, phi_in):\n",
    "    print(\"\\tphi_in\", phi_in)\n",
    "    print(\"\\tphi_out\", flux_surface.cylindrical_position(s = 2.0, theta = 0.0, phi = phi_in)[2])\n",
    "    \n",
    "print(\"For FluxSurfaceNormalExtended:\")\n",
    "check_same_phi(flux_surface_normal_extended, phi_in = 0.5)\n",
    "from jax_sbgeom.flux_surfaces import FluxSurfaceNormalExtendedNoPhi\n",
    "print(\"\\nFor FluxSurfaceNormalExtendedNoPhi:\")\n",
    "flux_surface_normal_extended_no_phi = FluxSurfaceNormalExtendedNoPhi.from_flux_surface(flux_surface)\n",
    "check_same_phi(flux_surface_normal_extended_no_phi, phi_in = 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1624817f",
   "metadata": {},
   "source": [
    "This sounds very good, but now they do not define the same surface anymore! We can plot both surfaces and see that they differ:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3c5186",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = create_figure_planes(ToroidalExtent.half_module(flux_surface_normal_extended))\n",
    "\n",
    "d_2_mesh_no_phi = jsb.flux_surfaces.mesh_surface(flux_surface_normal_extended_no_phi, s = 3.0, toroidal_extent=ToroidalExtent.half_module(flux_surface_normal_extended_no_phi), n_theta = 50, n_phi = 100)\n",
    "add_mesh_to_plotly(fig, lcfs_mesh, dict(opacity=0.3, color='pink', name = \"LCFS\"))\n",
    "add_mesh_to_plotly(fig, d_2_mesh_no_phi, dict(opacity=1.0, color='lightblue', name = \"d=2.0 no phi\"))\n",
    "add_mesh_to_plotly(fig, d_2_mesh, dict(opacity=0.6, color='red', name = \"d=2.0 with phi\"))\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b695aec",
   "metadata": {},
   "source": [
    "This might not be that important in some applications: it could be that the parametrisation is just some parametrisation and the precise meaning of the coordinate is not important. However, \n",
    "if you want a totally fixed distance from the LCFS, another option is to use Newton iterations. Then, we try to find a $\\varphi(\\phi_{in})$ such  that $\\mathbf{r}(\\varphi) = \\phi_{in}$ when extended with the normal vector. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24438469",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux_surface_normal_extended_constant_phi = jsb.flux_surfaces.FluxSurfaceNormalExtendedConstantPhi.from_flux_surface(flux_surface)\n",
    "print(\"For FluxSurfaceNormalExtendedConstantPhi:\")\n",
    "check_same_phi(flux_surface_normal_extended_constant_phi, phi_in = 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "addfaf3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = create_figure_planes(ToroidalExtent.half_module(flux_surface_normal_extended))\n",
    "\n",
    "d_2_mesh_constant_phi = jsb.flux_surfaces.mesh_surface(flux_surface_normal_extended_constant_phi, s = 3.0, toroidal_extent=ToroidalExtent.half_module(flux_surface_normal_extended_constant_phi), n_theta = 50, n_phi = 100)\n",
    "add_mesh_to_plotly(fig, lcfs_mesh, dict(opacity=0.3, color='pink', name = \"LCFS\"))\n",
    "add_mesh_to_plotly(fig, d_2_mesh_constant_phi, dict(opacity=1.0, color='lightblue', name = \"d=2.0 constant phi\"))\n",
    "add_mesh_to_plotly(fig, d_2_mesh, dict(opacity=0.6, color='red', name = \"d=2.0 with phi\"))\n",
    "\n",
    "fig\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af906d7e",
   "metadata": {},
   "source": [
    "However, this still has a price: $s>1$ is no longer a straight line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89bf3def",
   "metadata": {},
   "outputs": [],
   "source": [
    "positions_RZ = flux_surface_normal_extended_constant_phi.cylindrical_position(s = jnp.linspace(0,3.0, 50), theta = -0.7, phi = 0.5)\n",
    "fig = go.Figure(layout=go.Layout(width=600, height=500))\n",
    "plot_poloidal_slice(fig, flux_surface_normal_extended_constant_phi, phi = 0.5, s_values = [0.2, 0.5, 0.8, 1.0, 2.0], line_kwargs = {\"line\": dict(width=2)})\n",
    "fig.add_trace(go.Scatter(x = positions_RZ[:,0], y= positions_RZ[:,1], mode='lines', name=\"extension\", marker=dict(color='red', size=5)))\n",
    "fig.update_xaxes(scaleanchor=\"y\", scaleratio=1)\n",
    "fig "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4349d541",
   "metadata": {},
   "source": [
    "It depends on the application which extension is preferred: when ray-tracing from the LCFS, it can be very benificial to use the no-phi extension; then extensions are straight lines and thus a ray-trace is the same as an extension."
   ]
  }
 ],
 "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
}