Flux Surfaces
==============
We describe the basis ideas of the flux surfaces in :code:`jax-sbgeom` here. For more details, see the API reference and the examples.

Flux surfaces are a mapping:

.. math::
    (s, \theta, \varphi) \mapsto (R, Z, \phi)

The standard VMEC-like representation maps this as 

.. math::
    R(s, \theta, \varphi) &= \sum_{m,n} R_{mn}(s) \cos(m\theta - n\varphi)\\
    Z(s, \theta, \varphi) &= \sum_{m,n} Z_{mn}(s) \sin(m\theta - n\varphi)\\
    \phi(s, \theta, \varphi) &= \varphi

where the Fourier coefficients are functions of the flux surface label :math:`s`.
In VMEC, this is just a set of numbers, each corresponding to a fixed flux surface.

In :code:`jax-sbgeom`, this data is managed and stored in the abstract :class:`jax_sbgeom.flux_surfaces.FluxSurfaceBase` class.
The implementation :class:`jax_sbgeom.flux_surfaces.FluxSurface` then implements the above equations, up until the 
last-closed-flux-surface (LCFS), corresponding to :math:`s=1`.
From these equations, the tangent, normal and principal curvature vectors could be computed analytically,
but we just let JAX handle this using algorithmic differentation.

Distance beyond the last-closed flux surface
-------------------------------------------------

The plasma geometry can be extended by allowing :math:`s>1`. This is useful for 
describing e.g. blanket geometries.

Several different meanings can then be attached to coordinates with :math:`s>1`.

Normal Extension
~~~~~~~~~~~~~~~~~
We can just extend the geometry by using the normal vector on the LCFS, i.e. we can define for :math:`s>1`:

.. math::
    \mathbf{r}(s, \theta, \varphi) = \mathbf{r}(1, \theta, \varphi) + (s-1) \hat{\mathbf{n}}(1, \theta, \varphi)

This is implemented in the :class:`jax_sbgeom.flux_surfaces.FluxSurfaceNormalExtended` class.

Normal Extended No Phi
~~~~~~~~~~~~~~~~~~~~~~~~~~~
Using naive normal extension, we have that it no longer holds that :math:`\phi = \varphi` for :math:`s>1`. This can be inconvenient for some applications, e.g. meshing of the extended geometry
and Fourier transforming. The :class:`jax_sbgeom.flux_surfaces.FluxSurfaceNormalExtendedNoPhi` class implements a normal extension where the :math:`\phi = \varphi` constraint is still satisfied, by using the normal vector projected onto the :math:`R-Z` plane.

Normal Extended Constant Phi
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The no phi extension works well, but it loses the meaning of exact distance; the point could be located to another point on the LCFS closer than the distance given (i.e. :math:`s-1`).
Another option is to run a few Newton iterations to find a new :math:`\overline{\varphi}` such that :math:`\phi(1, \theta, \overline{\varphi}) = \varphi` for :math:`s>1`. 
This is implemented in the :class:`jax_sbgeom.flux_surfaces.FluxSurfaceNormalExtendedConstantPhi` class.

Fourier Extended
~~~~~~~~~~~~~~~~~~
Another option is to define a new :class:`jax_sbgeom.flux_surfaces.FluxSurfaceBase` object that represents the flux surface 
for :math:`s>1`. This is implemented in the :class:`jax_sbgeom.flux_surfaces.FluxSurfaceFourierExtended` class, which uses a Fourier representation for :math:`s>1` as well.


Meshing
-----------------
All of the above classes can be meshed using surface meshes (:func:`jax_sbgeom.flux_surfaces.mesh_surface`, :func:`jax_sbgeom.flux_surfaces.mesh_surfaces_closed`, :func:`jax_sbgeom.flux_surfaces.mesh_watertight_layers`) and volume meshes (:func:`jax_sbgeom.flux_surfaces.mesh_tetrahedra`).
See the API reference for details on the meshing functions: :mod:`jax_sbgeom.flux_surfaces.flux_surface_meshing`.


Converting between representations
--------------------------------------

Functions are included in the module :mod:`jax_sbgeom.flux_surfaces.convert_to_vmec` to convert between different representations, or to convert the parametrisation to an equal arclength.
Furthermore, given some thickness matrix (i.e. for :math:`\phi`, :math:`\theta` a distance to the LCFS), we can fit a flux surface to it.
This thickness matrix can be generated by :func:`jax_sbgeom.flux_surfaces.generate_thickness_matrix` using an external mesh.