jax_sbgeom.coils.base_coil module

class Coil

Bases: ABC

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

abstractmethod position(s)

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

abstractmethod tangent(s)

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

abstractmethod centre()

Centre of the coil

abstractmethod normal(s)

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

abstractmethod reverse_parametrisation()

coil_position(coil: Coil, s)

coil_tangent(coil: Coil, s)

coil_centre(coil: Coil)

coil_normal(coil: Coil, s)

coil_reverse_parametrisation(coil: Coil)

class FiniteSizeMethod

Bases: ABC

Abstract base class for finite size methods. All finite size method classes should inherit from this class and implement the abstract methods.

classmethod setup_from_coil(coil: Coil, *args)

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)

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. This base version results in all arguments being vmapped over.

Parameters:

• coil (Coil) – Coil object (batched over)

• args (tuple) – Additional arguments for the setup (specific to the FiniteSizeMethod). Possibly batched over, depending on the FiniteSizeMethod.

Returns:

Data needed to instantiate the FiniteSizeMethod (batched over)

Return type:

tuple

classmethod from_coil(coil: Coil, *args)

abstractmethod compute_radial_vector(coil: Coil, s: Array)

reverse_parametrisation()

‘ 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_finite_size_frame(coil: Coil, s: Array)

class FiniteSizeCoil(coil: Coil, finite_size_method: FiniteSizeMethod)

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)

position(s)

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

tangent(s)

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

centre()

Centre of the coil

normal()

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

reverse_parametrisation()

radial_vector(s)

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)

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)

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

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)

class FrenetSerretFrame

Bases: FiniteSizeMethod

Finite size method using Frenet-Serret frame.

compute_radial_vector(coil: Coil, s: Array)

class RadialVectorFrame(radial_vectors_i: Array)

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)

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)

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)

reverse_parametrisation()

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

class RotationMinimizedFrame(radial_vectors_i: Array)

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)

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)

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.

ensure_coil_rotation(coil: Coil, positive_rotation: bool)

Ensure that the coils in the coilset are ordered in the same rotation (positive is increasing angle when looking at R-Z cross section)

Parameters:

• coil (Coil) – Coil object to check

• positive_rotation (bool) – Whether the coil should have positive rotation (if False, negative rotation is enforced)

Returns

Coil : Coil with ensured rotation