jax_sbgeom.flux_surfaces package

class FluxSurface(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)

Bases: FluxSurfaceBase

cylindrical_position(s, theta, phi)

Cylindrical position on the flux surface as a function of (s, theta, phi)

Parameters:

• s (jnp.ndarray) – Surface index or normalized flux label

• theta (jnp.ndarray) – Poloidal angle(s)

• phi (jnp.ndarray) – Toroidal angle(s)

Returns:

Cylindrical position(s) on the flux surface [R, Z, phi]

Return type:

jnp.ndarray

cartesian_position(s, theta, phi)

Cartesian position on the flux surface as a function of (s, theta, phi)

Parameters:

• s (jnp.ndarray) – Surface index or normalized flux label

• theta (jnp.ndarray) – Poloidal angle(s)

• phi (jnp.ndarray) – Toroidal angle(s)

Returns:

Cartesian position(s) on the flux surface [x, y, z]

Return type:

jnp.ndarray

normal(s, theta, phi)

Normal vector on the flux surface as a function of (s, theta, phi)

Parameters:

• s (jnp.ndarray) – Surface index or normalized flux label

• theta (jnp.ndarray) – Poloidal angle(s)

• phi (jnp.ndarray) – Toroidal angle(s)

Returns:

Normal vector(s) on the flux surface

Return type:

jnp.ndarray

principal_curvatures(s, theta, phi)

Principal curvatures on the flux surface as a function of (s, theta, phi)

Parameters:

• s (jnp.ndarray) – Surface index or normalized flux label

• theta (jnp.ndarray) – Poloidal angle(s)

• phi (jnp.ndarray) – Toroidal angle(s)

Returns:

Principal curvatures(s) on the flux surface

Return type:

jnp.ndarray

class FluxSurfaceData(Rmnc: Array, Zmns: Array)

Bases: object

Rmnc: Array

Zmns: Array

classmethod from_rmnc_zmns_settings(Rmnc: Array, Zmns: Array, settings: FluxSurfaceSettings, make_normals_point_outwards: bool = True)

class FluxSurfaceModes(mpol_vector: Array, ntor_vector: Array)

Bases: object

mpol_vector: Array

ntor_vector: Array

classmethod from_settings(settings: FluxSurfaceSettings)

class FluxSurfaceSettings(mpol: int, ntor: int, nfp: int)

Bases: object

mpol: int

ntor: int

nfp: int

class ToroidalExtent(start: float, end: float)

Bases: object

Class representing a toroidal extent in phi.

Attributes:

startfloat

Starting toroidal angle (in radians)

endfloat

Ending toroidal angle (in radians)

start: float

end: float

classmethod half_module(flux_surface: FluxSurfaceBase, dphi=0.0)

Create a ToroidalExtent representing half a field period.

Parameters:

• flux_surface (FluxSurfaceBase) – The flux surface for which to create the toroidal extent.

• dphi (float) – An optional offset to add to both the start and end angles.

Returns:

The created ToroidalExtent.

Return type:

ToroidalExtent

classmethod full_module(flux_surface: FluxSurfaceBase, dphi=0.0)

Create a ToroidalExtent representing a full field period.

Parameters:

• flux_surface (FluxSurfaceBase) – The flux surface for which to create the toroidal extent.

• dphi (float) – An optional offset to add to both the start and end angles.

Returns:

The created ToroidalExtent.

Return type:

ToroidalExtent

classmethod full()

Create a ToroidalExtent representing the full toroidal angle [0, 2pi].

Returns:

The created ToroidalExtent.

Return type:

ToroidalExtent

full_angle()

Whether the toroidal extent represents a full 2pi angle. Used for creating closed meshes.

Returns:

True if the toroidal extent represents a full 2pi angle, False otherwise.

Return type:

bool

class FluxSurfaceBase(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)

Bases: ParametrisedSurface

Class representing a set of flux surfaces using a VMEC-like representation. Base abstract class that provides no implementation. The actual implementation is provided in the FluxSurface class, which inherits from this base class. This allows for different implementations of the same interface, such as interpolation-based or direct evaluation-based methods.

Attributes:

dataFluxSurfaceData

Data object containing the Fourier coefficients Rmnc and Zmns

modesFluxSurfaceModes

Modes object containing the mode vectors mpol_vector and ntor_vector

settingsFluxSurfaceSettings

Settings object containing parameters mpol, ntor, nfp

data: FluxSurfaceData = None

modes: FluxSurfaceModes = None

settings: FluxSurfaceSettings = None

classmethod from_hdf5(filename: str)

Load a FluxSurface from an VMEC-type HDF5 file.

Parameters:

filename (str) – Path to the HDF5 file.

Returns:

The loaded FluxSurface object.

Return type:

FluxSurface

classmethod from_flux_surface(flux_surface_base: FluxSurfaceBase)

Create a FluxSurface from another FluxSurface (copy constructor). Can be used to convert between subclasses of FluxSurface.

Parameters:

flux_surface_base (FluxSurfaceBase) – The input FluxSurface to copy from.

Returns:

The copied FluxSurface.

Return type:

FluxSurface

classmethod from_rmnc_zmns_settings(Rmnc: Array, Zmns: Array, settings: FluxSurfaceSettings, make_normals_point_outwards: bool = True)

Create a FluxSurface from Fourier coefficients and settings. Optionally, modify the coefficients such that normals point outwards (default=True)

Parameters:

• Rmnc (Array) – Array of radial Fourier coefficients. Shape (nsurf, nmodes)

• Zmns (Array) – Array of vertical Fourier coefficients. Shape (nsurf, nmodes)

• settings (FluxSurfaceSettings) – Settings object containing parameters mpol, ntor, nfp

• make_normals_point_outwards (bool) – Whether to modify the Fourier coefficients such that normals point outwards. Default is True.

Returns:

The created FluxSurface.

Return type:

FluxSurface

classmethod from_data_settings(data: FluxSurfaceData, settings: FluxSurfaceSettings)

Create a FluxSurface from FluxSurfaceData and FluxSurfaceSettings.

Parameters:

• data (FluxSurfaceData) – Data object containing the Fourier coefficients Rmnc and Zmns

• settings (FluxSurfaceSettings) – Settings object containing parameters mpol, ntor, nfp

Returns:

The created FluxSurface.

Return type:

FluxSurface

classmethod from_data_settings_full(data: FluxSurfaceData, settings: FluxSurfaceSettings)

Create a FluxSurface from FluxSurfaceData and FluxSurfaceSettings. In this function, the data is ensured to be 2D. This is required for some functions that interpolate between the different surfaces.

Parameters:

• data (FluxSurfaceData) – Data object containing the Fourier coefficients Rmnc and Zmns

• settings (FluxSurfaceSettings) – Settings object containing parameters mpol, ntor, nfp

Returns:

The created FluxSurface.

Return type:

FluxSurface

property nfp

class ParametrisedSurface

Bases: Module

• Internal *

Class representing a parametrised surface. This is used for the interpolation-based methods that require 2D data. The surface is parametrised by (s, theta, phi) where s is the surface index or normalized flux label, theta is the poloidal angle and phi is the toroidal angle.

abstractmethod cylindrical_position(s, theta, phi)

abstractmethod cartesian_position(s, theta, phi)

abstractmethod normal(s, theta, phi)

abstractmethod principal_curvatures(s, theta, phi)

class FluxSurfaceNormalExtended(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)

Bases: FluxSurfaceBase

Class representing a flux surface that is extended along the normal direction.

The extension is done such that: - For s <= 1.0, the original flux surface is used - For s > 1.0, the position is given by moving along the normal direction of the flux surface at s = 1.0

cartesian_position(s, theta, phi)

cylindrical_position(s, theta, phi)

normal(s, theta, phi)

principal_curvatures(s, theta, phi)

class FluxSurfaceNormalExtendedNoPhi(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)

Bases: FluxSurfaceBase

Class representing a flux surface that is extended along the normal direction, but with no toroidal (phi) component in the extension.

The extension is done such that: - For s <= 1.0, the original flux surface is used - For s > 1.0, the position is given by moving along the normal direction of the flux surface at s = 1.0, but with the toroidal component removed

This is useful for creating an extension label that preserves phi_in = phi_out but still extends in a straight line. However, the label does not have the meaning of ‘distance to the lcfs’ anymore, as the extension is not along the actual normal direction.

cartesian_position(s, theta, phi)

cylindrical_position(s, theta, phi)

normal(s, theta, phi)

principal_curvatures(s, theta, phi)

class FluxSurfaceNormalExtendedConstantPhi(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)

Bases: FluxSurfaceBase

Class representing a flux surface that is extended along the normal direction, but keeping the toroidal angle (phi) constant during the extension. The extension is done such that: - For s <= 1.0, the original flux surface is used - For s > 1.0, the position is given by moving along the normal direction of the flux surface at s = 1.0, but adjusting the toroidal angle to keep it constant

This is useful for creating an extension label that preserves phi_in = phi_out while retaining the meaning of ‘distance to the lcfs’, as the extension is still along the normal direction. However, the extension is no longer a straight line in 3D space.

cartesian_position(s, theta, phi)

cylindrical_position(s, theta, phi)

normal(s, theta, phi)

principal_curvatures(s, theta, phi)

class FluxSurfaceFourierExtended(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None, extension_flux_surface: FluxSurfaceBase = None)

Bases: FluxSurfaceBase

A flux surface that is extended using another flux surface defined in Fourier space. This does not necessarily have to have the same mpol & ntor as the inner flux surface.

The inner flux surface is used for s <= 1.0, and the extension flux surface is used for s > 1.0.

s = 1.0 corresponds to the LCFS of the inner surface s = 2.0 corresponds to the first surface of the extension surface etc.. s = n_extension + 1.0 corresponds to the last surface of the extension surface

Beyond that, the additional s is ignored.

extension_flux_surface: FluxSurfaceBase = None

classmethod from_flux_surface_and_extension(flux_surface: FluxSurfaceBase, extension_flux_surface: FluxSurfaceBase)

Create a FluxSurfaceFourierExtended from a base flux surface and an extension flux surface.

Parameters:

• flux_surface (FluxSurfaceBase) – Base flux surface to extend.

• extension_flux_surface (FluxSurfaceBase) – Extension flux surface to use for s > 1.0.

Return type:

FluxSurfaceFourierExtended

cartesian_position(s, theta, phi)

cylindrical_position(s, theta, phi)

normal(s, theta, phi)

principal_curvatures(s, theta, phi)

class FluxSurfaceExtendedDistanceMatrix(flux_surface_extended: FluxSurfaceBase, d_layers: Array)

Bases: ParametrisedSurface

flux_surface_extended: FluxSurfaceBase

Flux surface extenion used. Cannot be of type FluxSurface.

d_layers: Array

assumed to be shaped like [n_theta_sampled, n_phi_sampled]. s[0,0] is (0,0), s[-1-1] corresponds to (2pi, 2pi/nfp)

Type:

The distance matrices of each layer

cylindrical_position(s, theta, phi)

s_interp(s, theta, phi)

cartesian_position(s, theta, phi)

normal(s, theta, phi)

principal_curvatures(s, theta, phi)

mesh_surface(flux_surfaces: ParametrisedSurface, s: float, toroidal_extent: ToroidalExtent, n_theta: int, n_phi: int, normals_facing_outwards: bool = True)

Create a mesh of points on a flux surface at normalized radius s, with n_theta poloidal and n_phi toroidal points.

This cannot be jitted because the toroidal extent determines whether it is closed, which determines the number of triangles. Therefore, the size of the arrays is unknown at compile time (which cannot be jitted unless toroidal_extent is static, but this is inconvenient because then the function would recompile for every different extent).

This is therefore a convenience function only. Internal functions should not build on this function.

Parameters:

• flux_surfaces (ParametrisedSurface) – The parametrized surface object containing the parameterization.

• s (float) – The normalized radius of the flux surface to mesh.

• toroidal_extent (ToroidalExtent) – The toroidal extent of the mesh.

• n_theta (int) – The number of poloidal points.

• n_phi (int) – The number of toroidal points.

• normals_facing_outwards (bool) – Whether the normals should face outwards. Default is True.

Returns:

• positions (jnp.ndarray) – An array of shape (n_points, 3) containing the Cartesian coordinates of the mesh points.

• triangles (jnp.ndarray) – An array of shape (n_triangles, 3) containing the indices of the vertices for each triangle.

mesh_surfaces_closed(flux_surfaces: ParametrisedSurface, s_values_start: float, s_value_end: float, toroidal_extent: ToroidalExtent, n_theta: int, n_phi: int, n_cap: int)

Create a closed mesh of points on flux surfaces between normalized radius s_values_start and s_value_end, with n_theta poloidal and n_phi toroidal points.

This cannot be jitted because the toroidal extent determines whether it is closed, which determines the number of triangles. Therefore, the size of the arrays is unknown at compile time (which cannot be jitted unless toroidal_extent is static, but this is inconvenient because then the function would recompile for every different extent).

This is therefore a convenience function only. Internal functions should not build on this function.

Parameters:

• flux_surfaces (ParametrisedSurface) – The parametrized surface object containing the parameterization.

• s_values_start (float) – The starting normalized radius of the flux surfaces to mesh.

• s_value_end (float) – The ending normalized radius of the flux surfaces to mesh.

• include_axis (bool) – Whether to include the magnetic axis in the mesh.

• toroidal_extent (ToroidalExtent) – The toroidal extent of the mesh.

• n_theta (int) – The number of poloidal points.

• n_phi (int) – The number of toroidal points.

• n_cap (int) – The number of radial points in the caps if not full toroidal extent.

Returns:

• positions (jnp.ndarray) – An array of shape (n_points, 3) containing the Cartesian coordinates of the mesh points.

• triangles (jnp.ndarray) – An array of shape (n_triangles, 3) containing the indices of the vertices for each triangle.

mesh_watertight_layers(flux_surfaces: ParametrisedSurface, s_values: Array, toroidal_extent: ToroidalExtent, n_theta: int, n_phi: int)

Create a watertight mesh of points on flux surfaces at normalized radii s_values, with n_theta poloidal and n_phi toroidal points.

This cannot be jitted because the toroidal extent determines whether it is closed, which determines the number of triangles. Therefore, the size of the arrays is unknown at compile time (which cannot be jitted unless toroidal_extent is static, but this is inconvenient because then the function would recompile for every different extent).

This is therefore a convenience function only. Internal functions should not build on this function.

Parameters:

• flux_surfaces (ParametrisedSurface) – The parametrized surface object containing the parameterization.

• s_values (Array) – An array of normalized radii of the flux surfaces to mesh.

• toroidal_extent (ToroidalExtent) – The toroidal extent of the mesh.

• n_theta (int) – The number of poloidal points.

• n_phi (int) – The number of toroidal points.

Returns:

• positions (jnp.ndarray) – An array of shape (n_points, 3) containing the Cartesian coordinates of the mesh points.

• triangles (jnp.ndarray) – An array of shape (n_triangles, 3) containing the indices of the vertices for each triangle.

mesh_tetrahedra(flux_surfaces: ParametrisedSurface, s_values: Array, toroidal_extent: ToroidalExtent, n_theta: int, n_phi: int)

Create a tetrahedral mesh between layers of flux surfaces at normalized radii s_values, with n_theta poloidal and n_phi toroidal points.

This cannot be jitted because the toroidal extent determines whether it is closed, which determines the number of tetrahedra. Therefore, the size of the arrays is unknown at compile time (which cannot be jitted unless toroidal_extent is static, but this is inconvenient because then the function would recompile for every different extent).

This is therefore a convenience function only. Internal functions should not build on this function.

Parameters:

• flux_surfaces (ParametrisedSurface) – The flux surface object containing the parameterization.

• s_values (Array) – An array of normalized radii of the flux surfaces to mesh.

• toroidal_extent (ToroidalExtent) – The toroidal extent of the mesh.

• n_theta (int) – The number of poloidal points.

• n_phi (int) – The number of toroidal points.

Returns:

• positions (jnp.ndarray) – An array of shape (n_points, 3) containing the Cartesian coordinates of the mesh points.

• tetrahedra (jnp.ndarray) – An array of shape (n_tetrahedra, 4) containing the indices of the vertices for each tetrahedron.

create_fourier_representation(flux_surface: FluxSurfaceBase, s: Array, theta_grid: Array)

Create a Fourier representation of a flux surface at given (s, theta) grid points.

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to create the Fourier representation of.

• s (Array) – Radial coordinate(s) at which to sample the flux surface. If an array, must have the same shape as theta_grid.

• theta_grid (Array) – Grid of poloidal angles at which to sample the flux surface.

Return type:

Tuple[FluxSurfaceData, FluxSurfaceSettings]

Returns:

• flux_surface_data (FluxSurfaceData) – Fourier representation of the sampled flux surface.

• settings (FluxSurfaceSettings) – Settings of the Fourier representation (mpol, ntor, nfp).

convert_to_different_settings(fluxsurface: FluxSurfaceBase, settings_new: FluxSurfaceSettings)

Convert FluxSurface to a different (mpol, ntor) representation.

Note that this returns the same type as the input fluxsurface. However, if it is e.g. a FluxSurfaceFourierExtended, the extension data is not converted or used, so the return type will be only the base FluxSurface.

Parameters:

• fluxsurface (FluxSurfaceBase) – The flux surface to convert.

• settings_new (FluxSurfaceSettings) – The new Fourier settings (mpol, ntor).

Returns:

fluxsurface_new – New flux surface with Fourier coefficients in the new (mpol, ntor) representation. Same as type as input fluxsurface.

Return type:

FluxSurface

convert_to_equal_arclength(flux_surface: FluxSurfaceBase, n_theta: int, n_phi: int, n_theta_s_arclength: int)

Return type:

Tuple[FluxSurfaceData, FluxSurfaceSettings]

create_fourier_representation_d_interp(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int)

Create a Fourier representation of an extended flux surface with an interpolated extension distance. Can be batched over d: if d is a scalar or 2D array, a single flux surface is created. If d is a 1D or 3D array, multiple flux surfaces are created (batched).

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to extend using the distance function. Flux surface must be of type FluxSurfaceNormalExtendedNoPhi or FluxSurfaceNormalExtendedConstantPhi to ensure valid results (phi_in must be phi out for FFT)

• d (Array) – Distance function to extend the flux surface with. Assumed to be full module: i.e. phi in [0, 2pi/nfp], theta in [0, 2pi] (included endpoints) If d is a scalar or 2D array, a single flux surface is created. If d is a 1D or 3D array, multiple flux surfaces are created (batched).

• n_theta (int) – Number of poloidal points in the output Fourier representation.

• n_phi (int) – Number of toroidal points in the output Fourier representation.

Returns:

• flux_surface_data (FluxSurfaceData) – Fourier representation of the sampled flux surface.

• settings (FluxSurfaceSettings) – Settings of the Fourier representation (mpol, ntor, nfp).

create_fourier_representation_d_interp_equal_arclength(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int, n_theta_s_arclength: int)

Convenience function of create_fourier_representation_d_interp + convert_to_equal_arclength

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to extend using the distance function. Flux surface must be of type FluxSurfaceNormalExtendedNoPhi or FluxSurfaceNormalExtendedConstantPhi to ensure valid results (phi_in must be phi_out for FFT)

• d (Array) – Distance function to extend the flux surface with. Assumed to be full module: i.e. phi in [0, 2pi/nfp], theta in [0, 2pi] (included endpoints)

• n_theta (int) – Number of poloidal points in the output Fourier representation.

• n_phi (int) – Number of toroidal points in the output Fourier representation.

• n_theta_s_arclength (int) – Number of poloidal points to use for the arclength parametrization.

Returns:

• flux_surface_data (FluxSurfaceData) – Fourier representation of the sampled flux surface.

• settings (FluxSurfaceSettings) – Settings of the Fourier representation (mpol, ntor, nfp).

create_flux_surface_d_interp(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int, type_c: Type = <class 'jax_sbgeom.flux_surfaces.flux_surfaces_base.FluxSurface'>)

Convenience function of create_fourier_representation_d_interp + type_c.from_data_settings_full, returning a FluxSurface of given type.

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to extend using the distance function. Flux surface must be of type FluxSurfaceNormal

• d (Array) – Distance function to extend the flux surface with. Assumed to be full module: i.e. phi in [0, 2pi/nfp], theta in [0, 2pi] (included endpoints)

• n_theta (int) – Number of poloidal points in the output Fourier representation.

• n_phi (int) – Number of toroidal points in the output Fourier representation.

Returns:

flux_surface – Flux surface with Fourier representation.

Return type:

FluxSurface

create_flux_surface_d_interp_equal_arclength(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int, n_theta_s_arclength: int, type_c: Type = <class 'jax_sbgeom.flux_surfaces.flux_surfaces_base.FluxSurface'>)

Convenience function of create_fourier_representation_d_interp + convert_to_equal_arclength + type_c.from_data_settings_full, returning a FluxSurface of given type.

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to extend using the distance function. Flux surface must be of type FluxSurfaceNormalExtendedNoPhi or FluxSurfaceNormalExtendedConstantPhi to ensure valid results (phi_in must be phi_out for FFT)

• d (Array) – Distance function to extend the flux surface with. Assumed to be full module: i.e. phi in [0, 2pi/nfp], theta in [0, 2pi] (included endpoints)

• n_theta (int) – Number of poloidal points in the output Fourier representation.

• n_phi (int) – Number of toroidal points in the output Fourier representation.

• n_theta_s_arclength (int) – Number of poloidal points to use for the arclength parametrization.

Returns:

flux_surface – Flux surface with Fourier representation sampled on an equal arclength poloidal grid.

Return type:

FluxSurface

create_extended_flux_surface_d_interp(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int)

Creates a FluxSurfaceFourierExtended by extending a given flux surface using a distance function d and interpolating the distance function.

Convenience function of create_fourier_representation_d_interp + FluxSurface.from_data_settings_full + FluxSurfaceFourierExtended.from_flux_surface_and_extension, returning a FluxSurfaceFourierExtended.

Compared to create_flux_surface_d_interp, this function directly returns a FluxSurfaceFourierExtended.

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to extend using the distance function. Flux surface must be of type FluxSurfaceNormal

• d (Array) – Distance function to extend the flux surface with. Assumed to be full module: i.e. phi in [0, 2pi/nfp], theta in [0, 2pi] (included endpoints)

• n_theta (int) – Number of poloidal points in the output Fourier representation.

• n_phi (int) – Number of toroidal points in the output Fourier representation.

Returns:

flux_surface_extended – Extended flux surface with Fourier representation.

Return type:

FluxSurfaceFourierExtended

create_extended_flux_surface_d_interp_equal_arclength(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int, n_theta_s_arclength: int)

Creates a FluxSurfaceFourierExtended by extending a given flux surface using a distance function d, interpolating the distance function, and sampling on an equal arclength poloidal grid.

Convenience function of create_fourier_representation_d_interp + convert_to_equal_arclength + FluxSurface.from_data_settings_full + FluxSurfaceFourierExtended.from_flux_surface_and_extension, returning a FluxSurfaceFourierExtended.

Compared to create_flux_surface_d_interp_equal_arclength, this function directly returns a FluxSurfaceFourierExtended.

Parameters:

• flux_surface (FluxSurfaceBase) – Flux_Surface to extend using the distance function. Flux surface must be of type FluxSurfaceNormalExtendedNoPhi or FluxSurfaceNormalExtendedConstantPhi to ensure valid results (phi_in must be phi_out for FFT)

• d (Array) – Distance function to extend the flux surface with. Assumed to be full module: i.e. phi in [0, 2pi/nfp], theta in [0, 2pi] (included endpoints) If d is a scalar or 2D array, a single flux surface is created. If d is a 1D or 3D array, multiple flux surfaces are created (batched).

• n_theta (int) – Number of poloidal points in the output Fourier representation.

• n_phi (int) – Number of toroidal points in the output Fourier representation.

• n_theta_s_arclength (int) – Number of poloidal points to use for the arclength parametrization.

Returns:

flux_surface_extended – Extended flux surface with Fourier representation sampled on an equal arclength poloidal grid.

Return type:

FluxSurfaceFourierExtended

generate_thickness_matrix(flux_surface: FluxSurfaceNormalExtendedNoPhi, mesh, n_theta: int, n_phi: int)

Generate thickness matrix of an external mesh with respect to a no-phi extended flux surface.

Uses the internal raytracing utilities to compute the minimum distance from the flux surface to the mesh along the normal directions.

Parameters:

• flux_surface (FluxSurfaceNormalExtendedNoPhi) – Flux surface to compute thickness from.

• mesh (Tuple[jnp.ndarray, jnp.ndarray]) – Mesh of the external object (vertices, connectivity).

• n_theta (int) – Number of poloidal points.

• n_phi (int) – Number of toroidal points.

Returns:

• theta (jnp.ndarray [n_theta, n_phi]) – Poloidal angles of the thickness matrix.

• phi (jnp.ndarray [n_theta, n_phi]) – Toroidal angles of the thickness matrix.

• dmesh (jnp.ndarray [n_theta, n_phi]) – Thickness matrix values.

Submodules

• jax_sbgeom.flux_surfaces.convert_to_vmec module
  • convert_to_different_settings()
  • convert_to_equal_arclength()
  • mpol_ntor_from_ntheta_nphi()
  • create_fourier_representation()
  • create_fourier_representation_d_interp()
  • create_flux_surface_d_interp()
  • create_extended_flux_surface_d_interp()
  • create_fourier_representation_d_interp_equal_arclength()
  • create_flux_surface_d_interp_equal_arclength()
  • create_extended_flux_surface_d_interp_equal_arclength()
• jax_sbgeom.flux_surfaces.flux_surface_meshing module
  • mesh_poloidal_plane()
  • mesh_surface()
  • mesh_surfaces_closed()
  • mesh_tetrahedra()
  • mesh_watertight_layers()
• jax_sbgeom.flux_surfaces.flux_surfaces_base module
  • FluxSurfaceSettings
  • reverse_theta_total()
  • FluxSurfaceModes
  • FluxSurfaceData
  • ParametrisedSurface
  • FluxSurfaceBase
  • FluxSurface
  • make_2d_flux_surface()
  • ToroidalExtent
• jax_sbgeom.flux_surfaces.flux_surfaces_extended module
  • FluxSurfaceNormalExtended
  • FluxSurfaceNormalExtendedNoPhi
  • FluxSurfaceNormalExtendedConstantPhi
  • FluxSurfaceFourierExtended
  • FluxSurfaceExtendedDistanceMatrix
• jax_sbgeom.flux_surfaces.flux_surfaces_utilities module
  • generate_thickness_matrix()
  • convert_half_module_points_to_full_module()
  • convert_full_module_points_multiple_full_module()
  • convert_half_module_points_to_multiple_full_modules_mesh()
  • convert_f_half_mod_to_full_mod()