jax_sbgeom.flux_surfaces.convert_to_vmec module

convert_to_different_settings(fluxsurface: FluxSurfaceBase, settings_new: FluxSurfaceSettings)[source]

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

Tuple[FluxSurfaceData, FluxSurfaceSettings]

mpol_ntor_from_ntheta_nphi(n_theta: int, n_phi: int)[source]
Return type:

Tuple[int, int]

create_fourier_representation(flux_surface: FluxSurfaceBase, s: Array, theta_grid: Array)[source]

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

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

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

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_extended_flux_surface_d_interp(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int)[source]

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_fourier_representation_d_interp_equal_arclength(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int, n_theta_s_arclength: int)[source]

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

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_equal_arclength(flux_surface: FluxSurfaceBase, d: Array, n_theta: int, n_phi: int, n_theta_s_arclength: int)[source]

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