Basic Flux Surfaces

import jax_sbgeom as jsb
import jax.numpy as jnp
import jax 
jax.config.update("jax_enable_x64", True)
import pyvista
import plotly.graph_objects as go

We first import the data from some VMEC file (we use a proprietary file here, but any can be used after converting to netcdf4)

from jax_sbgeom.flux_surfaces import FluxSurface
flux_surface = FluxSurface.from_hdf5("helias5_vmec.nc4")

W0306 16:49:52.482424   24380 cuda_executor.cc:1802] GPU interconnect information not available: INTERNAL: NVML doesn't support extracting fabric info or NVLink is not used by the device.
W0306 16:49:52.485456   24329 cuda_executor.cc:1802] GPU interconnect information not available: INTERNAL: NVML doesn't support extracting fabric info or NVLink is not used by the device.

Then, we can obtain positions from it easily:

flux_surface.cartesian_position(s = 1.0,  theta = 0.2, phi = 0.3)

Array([22.79398104,  7.05100461,  0.7066726 ], dtype=float64)

This base class does not go beyond the LCFS (same output):

flux_surface.cartesian_position(s = 2.0,  theta = 0.2, phi = 0.3)

Array([22.79398104,  7.05100461,  0.7066726 ], dtype=float64)

We can plot a poloidal slice:

def plot_poloidal_slice(fig, flux_surface, phi, s_values, n_theta = 50, line_kwargs = {}):
    theta = jnp.linspace(0, 2 * jnp.pi, n_theta)
    s_mg, theta_mg = jnp.meshgrid(jnp.array(s_values), theta, indexing="ij")
    cylindrical_values = flux_surface.cylindrical_position(s_mg, theta_mg, phi)
    
    for i in range(cylindrical_values.shape[0]):
        line_kwargs['name'] = f"s = {s_values[i]:.2f}"
        fig.add_trace(go.Scatter(x=cylindrical_values[i, :,0], y=cylindrical_values[i, :,1], mode='lines', **line_kwargs))

fig = go.Figure(layout=go.Layout(width=600, height=500))
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)})
fig.update_xaxes(scaleanchor="y", scaleratio=1)
fig

This doesn’t have the capability yet to ge beyond the LCFS, as you can see in the figure above.

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. We can obtain this by a convenience method:

from jax_sbgeom.flux_surfaces import FluxSurfaceNormalExtended

flux_surface_normal_extended = FluxSurfaceNormalExtended.from_flux_surface(flux_surface)

fig = go.Figure(layout=go.Layout(width=600, height=500))
plot_poloidal_slice(fig,flux_surface_normal_extended, phi = 0.2, s_values = [0.1,0.5,0.9,1.0, 2.0])
fig.update_xaxes(scaleanchor="y", scaleratio=1)
fig

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.

To show this, we use the mesh functionality to mesh a surface beyond the LCFS and show that it intersects some constant-\(\phi\) planes:

from jax_sbgeom.flux_surfaces import ToroidalExtent
def Plot_Phi_Plane(phi, fig, rmin,rmax, zmin,zmax, alpha = 1, color = 'blue', **kwargs):
    '''
    Plots a simple plane
    '''
    x = [rmin * jnp.cos(phi), rmax * jnp.cos(phi), rmax * jnp.cos(phi), rmin * jnp.cos(phi), rmin * jnp.cos(phi)]
    y = [rmin * jnp.sin(phi), rmax * jnp.sin(phi), rmax * jnp.sin(phi), rmin * jnp.sin(phi), rmin * jnp.sin(phi)]
    z = [zmin, zmin, zmax, zmax, zmin]
        
    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)
    fig = fig.add_trace(go.Scatter3d(x=x, y=y, z=z, mode='lines', line=dict(color='black', width=3), showlegend=False), **kwargs)

def create_figure_planes(toroidal_extent):
    fig = go.Figure()
    fig.update_layout(autosize=False, width=600, height=500, margin=dict(l=0, r=0, b=0, t=0))
    Plot_Phi_Plane(phi = toroidal_extent.start, fig = fig, rmin = 16, rmax=28, zmin=-6, zmax=6, alpha = 0.7,  color = 'lightblue')
    Plot_Phi_Plane(phi = toroidal_extent.end  , fig = fig, rmin = 16, rmax=28, zmin=-6, zmax=6, alpha = 0.7,  color = 'lightblue')
    return fig

def add_mesh_to_plotly(fig, mesh, mesh_kwargs, trace_kwargs = {}):
    x,y,z = mesh[0][:,0], mesh[0][:,1], mesh[0][:,2]
    i,j,k = mesh[1][:,0], mesh[1][:,1], mesh[1][:,2]
    fig.add_trace(go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, showlegend=True, **mesh_kwargs), **trace_kwargs)

fig = create_figure_planes(ToroidalExtent.half_module(flux_surface_normal_extended))
lcfs_mesh = jsb.flux_surfaces.mesh_surface(flux_surface_normal_extended, s = 1.0, 
                                           toroidal_extent=ToroidalExtent.full_module(flux_surface_normal_extended, -0.25 * 2 * jnp.pi / flux_surface_normal_extended.nfp),
                                            n_theta = 50, 
                                            n_phi = 100)
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)
add_mesh_to_plotly(fig, lcfs_mesh, dict(opacity=0.3, color='pink', name = "LCFS"))
add_mesh_to_plotly(fig, d_2_mesh, dict(opacity=1.0, color='lightblue', name = "d=2.0"))

fig

Clearly, the surface intersects the half module planes.

To remedy this, we can project the normal vector on the RZ plane and use that to parametrise distance beyond the LCFS.

def check_same_phi(flux_surface, phi_in):
    print("\tphi_in", phi_in)
    print("\tphi_out", flux_surface.cylindrical_position(s = 2.0, theta = 0.0, phi = phi_in)[2])
    
print("For FluxSurfaceNormalExtended:")
check_same_phi(flux_surface_normal_extended, phi_in = 0.5)
from jax_sbgeom.flux_surfaces import FluxSurfaceNormalExtendedNoPhi
print("\nFor FluxSurfaceNormalExtendedNoPhi:")
flux_surface_normal_extended_no_phi = FluxSurfaceNormalExtendedNoPhi.from_flux_surface(flux_surface)
check_same_phi(flux_surface_normal_extended_no_phi, phi_in = 0.5)

For FluxSurfaceNormalExtended:
	phi_in 0.5

	phi_out 0.4972359404296961

For FluxSurfaceNormalExtendedNoPhi:
	phi_in 0.5

	phi_out 0.5

This sounds very good, but now they do not define the same surface anymore! We can plot both surfaces and see that they differ:

fig = create_figure_planes(ToroidalExtent.half_module(flux_surface_normal_extended))

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)
add_mesh_to_plotly(fig, lcfs_mesh, dict(opacity=0.3, color='pink', name = "LCFS"))
add_mesh_to_plotly(fig, d_2_mesh_no_phi, dict(opacity=1.0, color='lightblue', name = "d=2.0 no phi"))
add_mesh_to_plotly(fig, d_2_mesh, dict(opacity=0.6, color='red', name = "d=2.0 with phi"))

fig

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, 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.

flux_surface_normal_extended_constant_phi = jsb.flux_surfaces.FluxSurfaceNormalExtendedConstantPhi.from_flux_surface(flux_surface)
print("For FluxSurfaceNormalExtendedConstantPhi:")
check_same_phi(flux_surface_normal_extended_constant_phi, phi_in = 0.5)

For FluxSurfaceNormalExtendedConstantPhi:
	phi_in 0.5

	phi_out 0.5

fig = create_figure_planes(ToroidalExtent.half_module(flux_surface_normal_extended))

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)
add_mesh_to_plotly(fig, lcfs_mesh, dict(opacity=0.3, color='pink', name = "LCFS"))
add_mesh_to_plotly(fig, d_2_mesh_constant_phi, dict(opacity=1.0, color='lightblue', name = "d=2.0 constant phi"))
add_mesh_to_plotly(fig, d_2_mesh, dict(opacity=0.6, color='red', name = "d=2.0 with phi"))

fig

However, this still has a price: \(s>1\) is no longer a straight line:

positions_RZ = flux_surface_normal_extended_constant_phi.cylindrical_position(s = jnp.linspace(0,3.0, 50), theta = -0.7, phi = 0.5)
fig = go.Figure(layout=go.Layout(width=600, height=500))
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)})
fig.add_trace(go.Scatter(x = positions_RZ[:,0], y= positions_RZ[:,1], mode='lines', name="extension", marker=dict(color='red', size=5)))
fig.update_xaxes(scaleanchor="y", scaleratio=1)
fig 

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.