jax_sbgeom.coils package

class Coil[source]

Bases: ABC

Abstract base class for coils. All coil classes should inherit from this class and implement the abstract methods.

abstractmethod position(s)[source]

Position as a function of normalised arc length s (0 <= s < 1)

abstractmethod tangent(s)[source]

Tangent vector as a function of normalised arc length s (0 <= s < 1)

abstractmethod centre()[source]

Centre of the coil

abstractmethod normal(s)[source]

Normal vector as a function of normalised arc length s (0 <= s < 1)

abstractmethod reverse_parametrisation()[source]
class FiniteSizeCoil(coil: Coil, finite_size_method: FiniteSizeMethod)[source]

Bases: Coil

Class representing a coil with finite size defined by a FiniteSizeMethod.

Internally, is a coil but also has a coil as member. The coil member is used to compute the position, tangent, centre and normal. The FiniteSizeMethod member is used to compute the radial vector and finite size frame.

coil: Coil
finite_size_method: FiniteSizeMethod
classmethod from_coil(coil: Coil, finite_size_method: Type[FiniteSizeMethod], *args)[source]
position(s)[source]

Position as a function of normalised arc length s (0 <= s < 1)

tangent(s)[source]

Tangent vector as a function of normalised arc length s (0 <= s < 1)

centre()[source]

Centre of the coil

normal()[source]

Normal vector as a function of normalised arc length s (0 <= s < 1)

reverse_parametrisation()[source]
radial_vector(s)[source]

Radial vector along the coil as a function of arc length

Uses the FiniteSizeMethod to compute the radial vector Uses the finite size method to compute the radial vector

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Radial vector(s) along the coil

Return type:

jnp.ndarray […, 3]

finite_size_frame(s)[source]

Finite size frame along the coil as a function of arc length Uses the finite size method to compute the finite size frame.

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Finite size vector(s) along the coil

Return type:

jnp.ndarray […, 2, 3]

finite_size(s, width_radial: float, width_phi: float)[source]

Finite size along the coil as a function of arc length Uses the finite size method to compute the finite size.

Parameters:
  • s (jnp.ndarray) – Arc length(s) along the coil

  • width_radial (float) – Width in the first finite size direction

  • width_phi (float) – Width in the second finite size direction

Returns:

Finite size vector(s) along the coil

Return type:

jnp.ndarray […, 4, 3]

class CentroidFrame[source]

Bases: FiniteSizeMethod

Finite size method using centroid frame

Centroid frame is defined by the radial vector being the vector that is pointing from the coil centre to the coil position, projected onto the plane normal to the tangent

compute_radial_vector(coil: Coil, s: Array)[source]
class FrenetSerretFrame[source]

Bases: FiniteSizeMethod

Finite size method using Frenet-Serret frame.

compute_radial_vector(coil: Coil, s: Array)[source]
class RadialVectorFrame(radial_vectors_i: Array)[source]

Bases: FiniteSizeMethod

Finite size method using precomputed radial vectors It interpolates in between the radial vectors to compute the radial vector and finite size frame at arbitrary locations along the coil. The radial vectors are assumed to be given at uniform arc length intervals (endpoint not included).

radial_vectors_i: Array
classmethod setup_from_coil(coil: Coil, *args)[source]

Function to setup the FiniteSizeMethod from the coil object. This returns the data needed to instantiate the FiniteSizeMethod (in a dictionary)

This is useful to vmap over the coil object. If one would only have a from_coil, this instantiates the FiniteSizeMethod inside the vmap, which is not JIT compatible. Therefore, the idea is to create the data needed to instantiate the FiniteSizeMethod inside the vmap, and then instantiate only once.

Parameters:
  • coil (Coil) – Coil object

  • args (tuple) – Additional arguments for the setup (specific to the FiniteSizeMethod)

classmethod setup_from_coils(coil: Coil, *args)[source]

Since the radial vectors are precomputed, this just returns the same as setup_from_coil. There is no function called that does any computation so there is no need to call a vmapped function.

classmethod from_radial_vectors(radial_vectors: Array)[source]
reverse_parametrisation()[source]

‘ Reverse the finite size method. By default, does nothing. Some finite size methods may need to do something different (i.e. interpolated radial frames need to reverse their data)

Return type:

FiniteSizeMethod

compute_radial_vector(coil: Coil, s: Array)[source]
class RotationMinimizedFrame(radial_vectors_i: Array)[source]

Bases: RadialVectorFrame

Finite size method using rotation minimized frame. This is a subclass of RadialVectorFrame. The radial vectors are computed using the rotation minimizing frame algorithm.

classmethod setup_from_coil(coil: Coil, number_of_rmf_samples: int)[source]

Function to setup the FiniteSizeMethod from the coil object. This returns the data needed to instantiate the FiniteSizeMethod (in a dictionary)

This is useful to vmap over the coil object. If one would only have a from_coil, this instantiates the FiniteSizeMethod inside the vmap, which is not JIT compatible. Therefore, the idea is to create the data needed to instantiate the FiniteSizeMethod inside the vmap, and then instantiate only once.

Parameters:
  • coil (Coil) – Coil object

  • args (tuple) – Additional arguments for the setup (specific to the FiniteSizeMethod)

classmethod setup_from_coils(coil: Coil, *args)[source]

Vmap version of setup_from_coil This is required because there needs to be a choice made whether the extra arguments are vmapped over or not. Some coils, like RotationMinimizedFrame, require a static number, and this cannot be vmapped using the base vmap.

class CoilSet(coils: Coil)[source]

Bases: object

Class representing a set of coils. Includes methods for batch evaluation of coil properties. Including with the same coordinate or different coordinates for each coil.

Internally, the coils are stored as a batched Coil object. Therefore, no mixed representations are supported.

Example

>>> coil1 = DiscreteCoil.from_positions(jnp.stack([ jnp.array([1.0, 0.0, 0.0]), jnp.array( [0.0, 1.0, 0.0]), jnp.array([-1.0, 0.0, 0.0]), jnp.array([0.0, -1.0, 0.0]) ]))
>>> coil2 = DiscreteCoil.from_positions(jnp.stack([ jnp.array([2.0, 0.0, 0.0]), jnp.array( [0.0, 2.0, 0.0]), jnp.array([-2.0, 0.0, 0.0]), jnp.array([0.0, -2.0, 0.0]) ]))
>>> coilset = CoilSet.from_list([coil1, coil2])
>>> s = jnp.linspace(0, 1, 100)
>>> positions = coilset.position(s)  # shape (2, 100, 3)
>>> tangents = coilset.tangent(s)    # shape (2, 100, 3)
>>> normals = coilset.normal(s)      # shape (2, 100, 3)
>>> positions_diff = coilset.position_different_s(s[None, :].repeat(coilset.n_coils, axis=0))  # shape (2, 100, 3)
>>> coil1_copy = coilset[0]  # Get first coil in the coilset
coils: Coil

Coils in the coilset, stored as a batched Coil object. The batch dimension is the first dimension of all arrays in the Coil dataclass.

classmethod from_list(coils: List[Coil])[source]
centre()[source]
position(s)[source]
tangent(s)[source]
normal(s)[source]
position_different_s(s)[source]
tangent_different_s(s)[source]
normal_different_s(s)[source]
property n_coils

Number of coils in the coilset. Uses the batched data shape and therefore is static information (can be used in jax.jit compiled functions as static shape).

class FiniteSizeCoilSet(coils: Coil)[source]

Bases: CoilSet

Class representing a set of finite size coils. Includes methods for batch evaluation of coil properties. Including with the same coordinate or different coordinates for each coil. Internally, the coils are stored as a batched FiniteSizeCoil object. Therefore, no mixed representations are supported. Is a subclass of CoilSet, so all methods from CoilSet are also available.

Example

>>> coil1 = DiscreteCoil.from_positions(jnp.stack([ jnp.array([1.0, 0.0, 0.0]), jnp.array( [0.0, 1.0, 0.0]), jnp.array([-1.0, 0.0, 0.0]), jnp.array([0.0, -1.0, 0.0]) ]))
>>> coil2 = DiscreteCoil.from_positions(jnp.stack([ jnp.array([2.0, 0.0, 0.0]), jnp.array( [0.0, 2.0, 0.0]), jnp.array([-2.0, 0.0, 0.0]), jnp.array([0.0, -2.0, 0.0]) ]))
>>> coilset = FiniteSizeCoilSet.from_coils([coil1, coil2], CentroidFrame)
>>> s = jnp.linspace(0, 1, 100)
>>> radial_vectors = coilset.radial_vector(s)  # shape (2, 100, 3)
>>> frames = coilset.finite_size_frame(s)      # shape (2, 100, 3, 3)
>>> finite_sizes = coilset.finite_size(s, 0.1, 0.1)  # shape (2, 100, 4, 3)
>>> radial_vectors_diff = coilset.radial_vector_different_s(s[None, :].repeat(coilset.n_coils, axis=0))  # shape (2, 100, 3)
>>> coil1_copy = coilset[0]  # Get first coil in the coilset
classmethod from_list(coils: List[FiniteSizeCoil])[source]

Create a FiniteSizeCoilSet from a list of FiniteSizeCoil objects.

Parameters:

coils (List[FiniteSizeCoil]) – List of FiniteSizeCoil objects

Returns:

FiniteSizeCoilSet object

Return type:

FiniteSizeCoilSet

classmethod from_coils(coils: List[Coil], method: Type[FiniteSizeMethod], *args)[source]

Create a FiniteSizeCoilSet from a list of Coil objects and a FiniteSizeMethod. This method is applied to all coils in the list.

Parameters:
  • coils (List[Coil]) – List of Coil objects

  • method (Type[FiniteSizeMethod]) – FiniteSizeMethod to use for meshing

  • args (tuple) – Additional arguments for the FiniteSizeMethod setup

Returns:

FiniteSizeCoilSet object

Return type:

FiniteSizeCoilSet

classmethod from_coilset(coilset: CoilSet, method: Type[FiniteSizeMethod], *args)[source]
radial_vector(s)[source]
finite_size_frame(s)[source]
finite_size(s, width_radial: float, width_phi: float)[source]
radial_vector_different_s(s)[source]
finite_size_frame_different_s(s)[source]
finite_size_different_s(s, width_radial: float, width_phi: float)[source]
class FourierCoil(fourier_cos: Array, fourier_sin: Array, centre_i: Array)[source]

Bases: Coil

Class representing a coil defined by Fourier coefficients.

It uses the Fourier series to compute the position, tangent and normal along the coil.

x = centre_i[0] + sum_{n=1}^N [ fourier_cos[0, n] * cos(2 pi n s) + fourier_sin[0, n] * sin(2 pi n s) ] y = centre_i[1] + sum_{n=1}^N [ fourier_cos[1, n] * cos(2 pi n s) + fourier_sin[1, n] * sin(2 pi n s) ] z = centre_i[2] + sum_{n=1}^N [ fourier_cos[2, n] * cos(2 pi n s) + fourier_sin[2, n] * sin(2 pi n s) ]

For creating FourierCoil objects from discrete positions, use the function curve_to_fourier_coefficients’. The parametrisation can be converted to equal-arclength by using convert_fourier_coil_to_equal_arclength or convert_fourier_coilset_to_equal_arclength.

Parameters:
  • fourier_cos (Array) – Fourier cosine coefficients [N_modes, 3]

  • fourier_sin (Array) – Fourier sine coefficients [N_modes, 3]

  • centre_i (Array) – Centre of the coil [3]

fourier_cos: Array
fourier_sin: Array
centre_i: Array
position(s)[source]

Position along the coil as a function of arc length

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Cartesian position(s) along the coil

Return type:

jnp.ndarray

tangent(s)[source]

Tangent vector along the coil as a function of arc length

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Tangent vector(s) along the coil

Return type:

jnp.ndarray

normal(s)[source]

Normal vector along the coil as a function of arc length

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Normal vector(s) along the coil

Return type:

jnp.ndarray

centre()[source]

Centre of the coil

reverse_parametrisation()[source]
convert_to_fourier_coilset(coilset: CoilSet, n_modes: int = None)[source]

Convert a DiscreteCoil to a FourierCoil by computing Fourier coefficients from the discrete positions

Parameters:
  • coil (DiscreteCoil) – Discrete coil object

  • n_modes (int) – Number of Fourier modes to use. If None, uses (N+1)//2 modes where N is the number of discrete points in the coil.

Returns:

Fourier coil object

Return type:

FourierCoil

convert_to_fourier_coil(coil: DiscreteCoil, n_modes: int = None)[source]

Convert a DiscreteCoil to a FourierCoil by computing Fourier coefficients from the discrete positions

Parameters:
  • coil (DiscreteCoil) – Discrete coil object

  • n_modes (int) – Number of Fourier modes to use. If None, uses N/2 modes where N is the number of discrete points in the coil.

Returns:

Fourier coil object

Return type:

FourierCoil

convert_fourier_coil_to_equal_arclength(fourier_coil: FourierCoil, n_points_sample: int = None, n_points_desired: int = None, method: Literal['pchip', 'linear'] = 'pchip')[source]

Resample a Fourier coil to have points equally spaced in arc length.

Parameters:
  • fourier_coil (FourierCoil) – Fourier coil object with N modes

  • n_points_sample (int) – Number of points to sample the arc length inverse. If None, uses N*16 points.

  • n_points_desired (int) – Number of points to resample to. If None, uses N*4 points.

  • method (Literal['pchip', 'linear']) – Method to use for interpolating the arc length inverse. ‘pchip’ is significantly better while not increasing runtime. Most of the time is spent on computing the Fourier sums.

Returns:

Resampled positions along the coil

Return type:

FourierCoil [n_points_desired, 3]

convert_fourier_coilset_to_equal_arclength(fourier_coilset: CoilSet, n_points_sample: int = None, n_points_desired: int = None, method: Literal['pchip', 'linear'] = 'pchip')[source]

Resample a Fourier coilset to have points equally spaced in arc length.

Parameters

n_points_sampleint

Number of points to sample the arc length inverse. If None, uses N*16 points.

n_points_desiredint

Number of points to resample to. If None, uses N*4 points.

methodLiteral[‘pchip’, ‘linear’]

Method to use for interpolating the arc length inverse. ‘pchip’ is significantly better while not increasing runtime. Most of the time is spent on computing the Fourier sums.

Returns:

Resampled positions along the coil

Return type:

FourierCoil [n_coils, n_points_desired, 3]

class DiscreteCoil(positions: Array, _centre_i: Array)[source]

Bases: Coil

Class representing a coil defined by discrete positions and a centre. This centre is precomputed from the given positions. It does not feature in the computation of the coil positions; only the centre.

positions: Array
classmethod from_positions(positions: Array)[source]

Create a DiscreteCoil from discrete positions

Positions are assumed to be non-periodic (in other words, the first and last point are not equal).

Parameters:

positions (Array) – Cartesian positions of discrete coil points […, 3]

Returns:

DiscreteCoil object

Return type:

DiscreteCoil

position(s)[source]

Position along the coil as a function of arc length

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Cartesian position(s) along the coil

Return type:

jnp.ndarray

tangent(s)[source]

Tangent vector along the coil as a function of arc length

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Tangent vector(s) along the coil

Return type:

jnp.ndarray

centre()[source]

Centre of the coil

Returns:

Centre of the coil

Return type:

jnp.ndarray

normal(s)[source]

Normal vector along the coil as a function of arc length Not defined for DiscreteCoil

Parameters:

s (jnp.ndarray) – Arc length(s) along the coil

Returns:

Normal vector(s) along the coil (jnp.nan)

Return type:

jnp.ndarray

reverse_parametrisation()[source]

Reverse the parametrisation of the discrete coil. Reverses all points except the first (s=0 remains s=0).

mesh_coil_surface(coil: FiniteSizeCoil, n_s: int, width_radial: float, width_phi: float)[source]

Mesh the surface of a coil using a defined number of samples and coil width.

Parameters:
  • coil (FiniteSizeCoil) – Coil to mesh

  • n_s (int) – Number of samples along the coil

  • width_radial (float) – Radial width of the finite sized coil

  • width_phi (float) – Toroidal width of the finite sized coil

Returns:

  • jnp.ndarray – Vertices of the meshed coil surface

  • jnp.ndarray – Connectivity array of the meshed coil surface

mesh_coilset_surface(coils: FiniteSizeCoilSet, n_s: int, width_radial: float, width_phi: float)[source]

Mesh the surface of a coilset

The coils vertices are originally: [n_coils, n_s, 4, 3] (4 lines per coil)

The coils connectivity is originally: [n_coils, n_s, 4, 2, 3] (4 lines per coil, 2 triangles per quad) Both are reshaped to (-1,3) to facilate easier post processing.

Parameters:
  • coil (Coil) – Coil to mesh

  • n_s (int) – Number of samples along the coil

  • width_radial (float) – Radial width of the finite sized coil

  • width_phi (float) – Toroidal width of the finite sized coil

Returns:

  • jnp.ndarray – Vertices of the meshed coil surface

  • jnp.ndarray – Connectivity array of the meshed coil surface

mesh_coil_volumetric(coil: FiniteSizeCoil, n_s: int, n_grid: int, width_radial: float, width_phi: float)[source]

Mesh the volume of a coil using a defined number of samples and coil width.

Parameters:
  • coil (FiniteSizeCoil) – Coil to mesh

  • n_s (int) – Number of samples along the coil

  • n_grid (int) – Number of grid points in the radial and toroidal directions

  • width_radial (float) – Radial width of the finite sized coil

  • width_phi (float) – Toroidal width of the finite sized coil

Returns:

  • jnp.ndarray – Vertices of the meshed coil volume

  • jnp.ndarray – Connectivity array of the meshed coil volume (tetrahedra)

create_coil_winding_surface(coilset: CoilSet, n_points_per_coil: int, n_points_phi: int, surface_type: Literal['spline', 'fourier', 'direct'] = 'spline')[source]

Create a coil winding surface from a CoilSet. The CoilSet is first ordered in phi and ensured to have positive orientation.

Parameters:
  • coilset (CoilSet) – CoilSet containing the coils to optimize.

  • n_points_per_coil (int) – Number of points per coil in the output mesh.

  • n_points_phi (int) – Number of points in the toroidal direction if needed.

  • surface_type (Literal['spline', 'fourier', 'direct']) – Method to create the surface: - “spline” uses a 3D periodic spline on each toroidal line, - “fourier” uses a fourier transformation on each toroidal line - “direct” meshes directly between the coils (no intermediate points)

Returns:

  • positions (jnp.ndarray [n_points, 3]) – Positions of the coil winding surface mesh points.

  • connectivity (jnp.ndarray [n_faces, 3]) – Connectivity of the coil winding surface mesh.

create_optimized_coil_winding_surface(coilset: CoilSet, n_points_per_coil: int, n_points_phi: int, surface_type: Literal['spline', 'fourier', 'direct'] = 'spline', uniformity_penalty: float = 1.0, repulsive_penalty: float = 0.1, n_samples_per_coil_opt: int = 100, optimization_settings=OptimizationSettings(max_iterations=100, tolerance=0.0001))[source]

Create an optimized coil winding surface mesh from a CoilSet. The CoilSet is first ordered in phi and ensured to have positive orientation.

Parameters:
  • coilset (CoilSet) – CoilSet containing the coils to optimize.

  • n_points_per_coil (int) – Number of points per coil in the output mesh.

  • n_points_phi (int) – Number of points in the toroidal direction if needed.

  • surface_type (Literal['spline', 'fourier', 'direct']) – Method to create the surface: - “spline” uses a 3D periodic spline on each toroidal line, - “fourier” uses a fourier transformation on each toroidal line - “direct” meshes directly between the coils (no intermediate points)

  • uniformity_penalty (float) – Weight of the uniformity loss.

  • repulsive_penalty (float) – Weight of the repulsion loss.

  • n_samples_per_coil_opt (int) – Number of sample points per coil for the optimization.

  • optimization_settings (OptimizationSettings) – Settings for the optimization process.

Returns:

  • positions (jnp.ndarray [n_points, 3]) – Positions of the coil winding surface mesh points.

  • connectivity (jnp.ndarray [n_faces, 3]) – Connectivity of the coil winding surface mesh.

calculate_normals_from_closest_point_on_mesh(coil: Coil | CoilSet, external_mesh, n_coil_samples: int)[source]

Calculate normals at coil positions by finding the closest points on an external mesh and using that normal.

Parameters:
  • coil (Union[Coil, CoilSet]) – Coil or CoilSet containing the coil(s).

  • external_mesh (jax_sbgeom.jax_utils.mesh.Mesh) – External mesh to find closest points on.

  • n_coil_samples (int) – Number of samples along the coil(s).

Returns:

  • positions (jnp.ndarray [n_coils, n_coil_samples, 3]) – Positions along the coil(s).

  • normals (jnp.ndarray [n_coils, n_coil_samples, 3]) – Normals at the coil positions.

Submodules