jax_sbgeom.flux_surfaces.flux_surfaces_base module

class FluxSurfaceSettings(mpol: int, ntor: int, nfp: int)[source]

Bases: object

mpol: int
ntor: int
nfp: int
reverse_theta_total(m_vec, n_vec, coeff_vec, cosine_sign: bool)

Vectorized version of _reverse_theta_single. Takes similar arguments as _reverse_theta_single but with additional array axes over which _reverse_theta_single is mapped.

Original documentation:

  • Internal *

Changes the Fourier coefficients such that theta is replaced by -theta.

m_vecjnp.ndarray

Array of poloidal mode numbers.

n_vecjnp.ndarray

Array of toroidal mode numbers.

coeff_vecjnp.ndarray

Array of Fourier coefficients. (Rmnc or Zmns)

cosine_signbool

If True, the coefficients correspond to cosine terms. If False, they correspond to sine terms.

new_coeff_vecjnp.ndarray

The modified Fourier coefficients after reversing theta.

class FluxSurfaceModes(mpol_vector: Array, ntor_vector: Array)[source]

Bases: object

mpol_vector: Array
ntor_vector: Array
classmethod from_settings(settings: FluxSurfaceSettings)[source]
class FluxSurfaceData(Rmnc: Array, Zmns: Array)[source]

Bases: object

Rmnc: Array
Zmns: Array
classmethod from_rmnc_zmns_settings(Rmnc: Array, Zmns: Array, settings: FluxSurfaceSettings, make_normals_point_outwards: bool = True)[source]
class ParametrisedSurface[source]

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)[source]
abstractmethod cartesian_position(s, theta, phi)[source]
abstractmethod normal(s, theta, phi)[source]
abstractmethod principal_curvatures(s, theta, phi)[source]
class FluxSurfaceBase(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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 FluxSurface(data: FluxSurfaceData = None, modes: FluxSurfaceModes = None, settings: FluxSurfaceSettings = None)[source]

Bases: FluxSurfaceBase

cylindrical_position(s, theta, phi)[source]

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)[source]

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

volume(s=1.0)[source]

Volume enclosed by the flux surface at s (divergence theorem, trapezoid rule over one field period).

Parameters:

s (float) – Surface index or normalized flux label (default: last closed flux surface)

Returns:

Enclosed volume

Return type:

float

area(s=1.0, oversampling: int = 4)[source]

Area of the flux surface at s.

Parameters:
  • s (float) – Surface index or normalized flux label (default: last closed flux surface)

  • oversampling (int) – Quadrature sampling relative to the volume integration (default 4, relative error ~1e-12 on a stellarator); the error decreases geometrically with it.

Returns:

Surface area

Return type:

float

mean_cross_section_area(s=1.0)[source]

phi-averaged area of the poloidal cross section at constant cylindrical angle phi (VMEC / simsopt definition).

A(phi) is the closed integral of R dZ over theta in the (R, Z) plane; the result is its average over phi. Both integrals are periodic trapezoid rules over the full theta range and one field period in phi. The integrand is band limited (twice the mode numbers), so the sampling used (6 * mpol + 1, 6 * ntor + 1) is exact to round-off.

Parameters:

s (float) – Surface index or normalized flux label (default: last closed flux surface)

Returns:

Mean cross sectional area

Return type:

float

minor_radius(s=1.0)[source]

Effective minor radius sqrt(mean_cross_section_area / pi) (VMEC Aminor).

Parameters:

s (float) – Surface index or normalized flux label (default: last closed flux surface)

Returns:

Minor radius

Return type:

float

major_radius(s=1.0)[source]

Effective major radius volume / (2 pi^2 minor_radius^2) (VMEC Rmajor).

Parameters:

s (float) – Surface index or normalized flux label (default: last closed flux surface)

Returns:

Major radius

Return type:

float

aspect_ratio(s=1.0)[source]

Aspect ratio major_radius / minor_radius.

Parameters:

s (float) – Surface index or normalized flux label (default: last closed flux surface)

Returns:

Aspect ratio

Return type:

float

normal(s, theta, phi)[source]

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)[source]

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

make_2d_flux_surface(fs: FluxSurfaceBase)[source]

Convert a FluxSurface with 1D data to a FluxSurface with 2D data by adding a dummy surface dimension. This allows you to use the interpolation-based methods that require 2D data. Interpolation on one surface will just return the single surface.

Parameters:

fs (FluxSurfaceBase) – The input FluxSurfaceBase with 1D data.

Returns:

The output FluxSurfaceBase with 2D data.

Return type:

FluxSurfaceBase

class ToroidalExtent(start: float, end: float)[source]

Bases: Module

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)[source]

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)[source]

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()[source]

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

Returns:

The created ToroidalExtent.

Return type:

ToroidalExtent

full_angle()[source]

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 ToroidalExtentStatic(start: float, end: float, full_angle_bool: bool)[source]

Bases: ToroidalExtent

Class extending a ToroidalExtent by having a static flag indicating whether it is a full angle.

For normal ToroidalExtent, the full angle is determined by checking whether end - start is close to 2pi. This can lead to issues when using it under jit for meshing; this will lead to different sized arrays depending on runtime values, which is not supported.

full_angle_bool: bool

Whether the toroidal extent represents a full 2 pi angle.

classmethod half_module(flux_surface: FluxSurfaceBase, dphi=0.0)[source]

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)[source]

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()[source]

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

Returns:

The created ToroidalExtent.

Return type:

ToroidalExtent

full_angle()[source]

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