API Reference
Tracing
raytrax.trace(magnetic_configuration: MagneticConfiguration, radial_profiles: RadialProfiles, beam: Beam, trim: bool = True, settings: TracerSettings = TracerSettings()) -> TraceResult
Trace a single beam through the plasma.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
magnetic_configuration
|
MagneticConfiguration
|
Magnetic configuration with gridded data |
required |
radial_profiles
|
RadialProfiles
|
Radial profiles of plasma parameters |
required |
beam
|
Beam
|
Beam initial conditions |
required |
trim
|
bool
|
If |
True
|
settings
|
TracerSettings
|
ODE solver settings (tolerances, step sizes, termination
thresholds). Defaults to :class: |
TracerSettings()
|
Returns:
| Type | Description |
|---|---|
TraceResult
|
TraceResult with beam profile and radial deposition profile. |
Input Types
raytrax.equilibrium.interpolate.MagneticConfiguration
dataclass
Bases: SafetensorsMixin
Magnetic configuration and geometry on a cylindrical grid.
Contains the magnetic field \(\boldsymbol{B}\) and normalized effective radius \(\rho\) on a 3D cylindrical grid (\(R\), \(\phi\), \(Z\)), along with volume information for computing deposition profiles.
raytrax.equilibrium.interpolate.MagneticConfiguration.dvolume_drho: jt.Float[jax.Array, ' nrho_1d']
instance-attribute
Volume derivative \(dV/d\rho\) on the 1D radial grid.
raytrax.equilibrium.interpolate.MagneticConfiguration.is_axisymmetric: bool = dataclass_field(default=False, metadata={'static': True})
class-attribute
instance-attribute
Whether the configuration is axisymmetric (tokamak) or 3D (stellarator).
raytrax.equilibrium.interpolate.MagneticConfiguration.is_stellarator_symmetric: bool = dataclass_field(metadata={'static': True})
class-attribute
instance-attribute
Whether the configuration has stellarator symmetry.
raytrax.equilibrium.interpolate.MagneticConfiguration.magnetic_field: jt.Float[jax.Array, 'npoints 3']
instance-attribute
The magnetic field \((B_R, B_\phi, B_Z)\) in cylindrical components at each grid point.
raytrax.equilibrium.interpolate.MagneticConfiguration.nfp: int = dataclass_field(metadata={'static': True})
class-attribute
instance-attribute
Number of field periods (toroidal periodicity).
raytrax.equilibrium.interpolate.MagneticConfiguration.rho: jt.Float[jax.Array, ' npoints']
instance-attribute
The normalized effective minor radius at each point on the interpolation grid.
raytrax.equilibrium.interpolate.MagneticConfiguration.rho_1d: jt.Float[jax.Array, ' nrho_1d']
instance-attribute
1D radial grid for volume derivative.
raytrax.equilibrium.interpolate.MagneticConfiguration.rphiz: jt.Float[jax.Array, 'npoints 3']
instance-attribute
The (\(R\), \(\phi\), \(Z\)) coordinates of the points on the interpolation grid.
raytrax.equilibrium.interpolate.MagneticConfiguration.from_vmec_wout(equilibrium: WoutLike, magnetic_field_scale: float = 1.0, grid: VmecGridResolution | None = None) -> MagneticConfiguration
classmethod
Create a MagneticConfiguration from a VMEC++ equilibrium.
Optionally, apply a uniform scale factor to the magnetic field magnitude with respect to the equilibrium data.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
equilibrium
|
WoutLike
|
an MHD equilibrium compatible with |
required |
magnetic_field_scale
|
float
|
Factor to multiply all magnetic field values by. |
1.0
|
grid
|
VmecGridResolution | None
|
Grid resolution settings. Defaults to
:class: |
None
|
Returns:
| Type | Description |
|---|---|
MagneticConfiguration
|
A MagneticConfiguration object containing interpolation data. |
raytrax.equilibrium.interpolate.MagneticConfiguration.to_pyvista_grid() -> pv.StructuredGrid
Convert to a PyVista StructuredGrid in Cartesian coordinates.
Point arrays written:
rho: normalised effective minor radius (NaN outside LCFS)B: magnetic field vector :math:(B_R, B_\phi, B_Z)in TabsB: magnetic field magnitude :math:|\boldsymbol{B}|in T
Requires the pyvista package.
Returns:
| Name | Type | Description |
|---|---|---|
A |
StructuredGrid
|
class: |
StructuredGrid
|
directly via |
raytrax.equilibrium.interpolate.CylindricalGridResolution
dataclass
Output cylindrical grid resolution, shared by all equilibrium importers.
Defines the \((R, \phi, Z)\) cylindrical grid on which the magnetic
configuration is stored and the ray tracer interpolates. For axisymmetric
(tokamak) configurations n_phi is ignored.
Attributes:
| Name | Type | Description |
|---|---|---|
n_r |
int
|
Number of grid points along the major radius \(R\). |
n_z |
int
|
Number of grid points along the vertical coordinate \(Z\). |
n_phi |
int
|
Number of toroidal planes \(\phi\). Ignored for axisymmetric configurations. |
n_rho_profile |
int
|
Number of radial points for the 1-D \(dV/d\rho\) profile. |
raytrax.equilibrium.interpolate.VmecGridResolution
dataclass
Grid resolution for VMEC-based equilibrium imports.
Combines the shared :class:CylindricalGridResolution output grid with
VMEC-specific intermediate flux-coordinate sampling parameters.
The VMEC pipeline first evaluates Fourier series on an intermediate
curvilinear \((\rho, \theta, \phi)\) grid, then scatter-interpolates the
result onto the cylindrical output grid. The intermediate grid parameters
n_rho and n_theta control the accuracy of that Fourier evaluation
step; they are invisible to the ray tracer.
Attributes:
| Name | Type | Description |
|---|---|---|
cylindrical |
CylindricalGridResolution
|
Output cylindrical grid shared with all other importers. |
n_rho |
int
|
Number of radial (flux-surface) points on the intermediate VMEC curvilinear grid. |
n_theta |
int
|
Number of poloidal points on the intermediate VMEC curvilinear grid. |
raytrax.types.Beam
dataclass
Beam parameter inputs for tracing.
raytrax.types.Beam.direction: jt.Float[jax.Array, 3]
instance-attribute
The starting direction of the beam in cartesian coordinates. Must be a unit vector.
raytrax.types.Beam.frequency: jt.Float[jax.Array, '']
instance-attribute
The frequency of the beam in Hz (not GHz!).
raytrax.types.Beam.mode: Literal['X', 'O'] = dataclass_field(metadata={'static': True})
class-attribute
instance-attribute
The polarization mode of the beam, either "X" for extraordinary or "O" for ordinary mode.
raytrax.types.Beam.position: jt.Float[jax.Array, 3]
instance-attribute
The starting position of the beam in cartesian coordinates.
raytrax.types.Beam.power: float
instance-attribute
The initial power of the beam in W (not MW!).
raytrax.types.RadialProfiles
dataclass
Radial profiles of electron density and temperature.
Defines the plasma parameters (electron density, temperature) as functions of the normalized effective radius \(\rho\).
raytrax.types.RadialProfiles.electron_density: jt.Float[jax.Array, ' nrho']
instance-attribute
The electron density profile in units of \(10^{20}\) m\(^{-3}\).
raytrax.types.RadialProfiles.electron_temperature: jt.Float[jax.Array, ' nrho']
instance-attribute
The electron temperature profile in keV.
raytrax.types.RadialProfiles.rho: jt.Float[jax.Array, ' nrho']
instance-attribute
The normalized effective minor radius grid.
raytrax.types.TracerSettings
dataclass
Settings for the ray tracing ODE solver.
All float fields are JAX-traceable (changing them does not trigger
recompilation). Only max_steps in the solver — which controls static
buffer sizes — would require recompilation, and is intentionally omitted
here for that reason.
raytrax.types.TracerSettings.absolute_tolerance: float = 1e-06
class-attribute
instance-attribute
Absolute tolerance of the adaptive step-size controller.
raytrax.types.TracerSettings.max_arc_length: float = 20.0
class-attribute
instance-attribute
Maximum arc length to integrate in metres before the solver gives up.
raytrax.types.TracerSettings.max_step_size: float = 0.05
class-attribute
instance-attribute
Maximum ODE step size in metres (arc length). Decreasing this improves accuracy at the cost of more steps.
raytrax.types.TracerSettings.relative_tolerance: float = 0.0001
class-attribute
instance-attribute
Relative tolerance of the adaptive step-size controller.
Output Types
raytrax.types.TraceResult
dataclass
The result of a ray tracing calculation.
raytrax.types.TraceResult.absorbed_power: jt.Float[jax.Array, '']
instance-attribute
Total power absorbed by the plasma in W.
raytrax.types.TraceResult.absorbed_power_fraction: jt.Float[jax.Array, '']
instance-attribute
Fraction of input beam power absorbed by the plasma, i.e. \(1 - e^{-\tau}\).
raytrax.types.TraceResult.beam_profile: BeamProfile
instance-attribute
The traced beam profile in real space.
raytrax.types.TraceResult.deposition_rho_mean: jt.Float[jax.Array, '']
instance-attribute
Flux-surface-volume-weighted mean normalised radius \(\langle\rho\rangle\) of power deposition.
raytrax.types.TraceResult.deposition_rho_std: jt.Float[jax.Array, '']
instance-attribute
Flux-surface-volume-weighted standard deviation of \(\rho\) for power deposition.
raytrax.types.TraceResult.optical_depth: jt.Float[jax.Array, '']
instance-attribute
Total optical depth \(\tau\) accumulated along the ray.
raytrax.types.TraceResult.radial_profile: RadialProfile
instance-attribute
The radial deposition profile.
raytrax.types.BeamProfile
dataclass
Beam profile in real space resulting from tracing.
raytrax.types.BeamProfile.absorption_coefficient: jt.Float[jax.Array, ' npoints']
instance-attribute
The absorption coefficient along the beam.
raytrax.types.BeamProfile.arc_length: jt.Float[jax.Array, ' npoints']
instance-attribute
The arc length along the beam.
raytrax.types.BeamProfile.electron_density: jt.Float[jax.Array, ' npoints']
instance-attribute
The electron density along the beam in units of \(10^{20}\) m$^{-3}.
raytrax.types.BeamProfile.electron_temperature: jt.Float[jax.Array, ' npoints']
instance-attribute
The electron temperature along the beam in keV.
raytrax.types.BeamProfile.linear_power_density: jt.Float[jax.Array, ' npoints']
instance-attribute
The linear power density along the beam.
raytrax.types.BeamProfile.magnetic_field: jt.Float[jax.Array, 'npoints 3']
instance-attribute
The magnetic field vector along the beam in T.
raytrax.types.BeamProfile.normalized_effective_radius: jt.Float[jax.Array, ' npoints']
instance-attribute
The normalized effective minor radius \(\rho\) along the beam.
raytrax.types.BeamProfile.optical_depth: jt.Float[jax.Array, ' npoints']
instance-attribute
The optical depth along the beam.
raytrax.types.BeamProfile.position: jt.Float[jax.Array, 'npoints 3']
instance-attribute
The position of the beam in cartesian coordinates.
raytrax.types.BeamProfile.refractive_index: jt.Float[jax.Array, 'npoints 3']
instance-attribute
The refractive index vector at each point along the beam.
raytrax.types.RadialProfile
dataclass
Beam profile in radial coordinates.
raytrax.types.RadialProfile.rho: jt.Float[jax.Array, ' npoints']
instance-attribute
The normalized effective minor radius.
raytrax.types.RadialProfile.volumetric_power_density: jt.Float[jax.Array, ' npoints']
instance-attribute
The volumetric power density in W/m³.
Physics
raytrax.physics.hamiltonian.hamiltonian(position: Float[jax.Array, 3], refractive_index: Float[jax.Array, 3], magnetic_field_interpolator: Callable[[Float[jax.Array, 3]], Float[jax.Array, 3]], rho_interpolator: Callable[[Float[jax.Array, 3]], Float[jax.Array, '']], electron_density_profile_interpolator: Callable[[Float[jax.Array, '']], Float[jax.Array, '']], frequency: Float[jax.Array, ''], mode: Literal['X', 'O']) -> tuple[Float[jax.Array, ''], HamiltonianAux]
Ray-tracing Hamiltonian \(\mathcal{H}(\boldsymbol{r}, \boldsymbol{n}) = |\boldsymbol{n}|^2 - n_\mathrm{AH}^2(\boldsymbol{r}, \boldsymbol{n})\).
Rays propagate along level sets \(\mathcal{H} = 0\). In vacuum (\(n_e < 10^{-6} \times 10^{20}\,\mathrm{m}^{-3}\)) the trivial form \(\mathcal{H} = |\boldsymbol{n}|^2 - 1\) is used; elsewhere the cold-plasma Appleton-Hartree dispersion relation selects the O- or X-mode refractive index squared.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
position
|
Float[Array, 3]
|
Cartesian position vector \(\boldsymbol{r}\) in metres. |
required |
refractive_index
|
Float[Array, 3]
|
Refractive index vector \(\boldsymbol{n} = c\boldsymbol{k}/\omega\). |
required |
magnetic_field_interpolator
|
Callable[[Float[Array, 3]], Float[Array, 3]]
|
Callable mapping a Cartesian position to the magnetic field vector \(\boldsymbol{B}\) in Tesla. |
required |
rho_interpolator
|
Callable[[Float[Array, 3]], Float[Array, '']]
|
Callable mapping a Cartesian position to the normalised effective radius \(\rho \in [0, 1]\). |
required |
electron_density_profile_interpolator
|
Callable[[Float[Array, '']], Float[Array, '']]
|
Callable mapping \(\rho\) to the electron density in \(10^{20}\,\mathrm{m}^{-3}\). |
required |
frequency
|
Float[Array, '']
|
Wave frequency \(f\) in Hz. |
required |
mode
|
Literal['X', 'O']
|
Polarisation mode — |
required |
Returns:
| Type | Description |
|---|---|
tuple[Float[Array, ''], HamiltonianAux]
|
Scalar value of the Hamiltonian. Zero on the dispersion surface. |
raytrax.physics.hamiltonian.hamiltonian_gradients = jax.grad(hamiltonian, argnums=(0, 1), has_aux=True)
module-attribute
Compute both Hamiltonian gradients \((\partial \mathcal{H}/\partial \boldsymbol{r},\, \partial \mathcal{H}/\partial \boldsymbol{n})\) in a single backward pass.
Signature mirrors hamiltonian: (position, refractive_index,
magnetic_field_interpolator, rho_interpolator,
electron_density_profile_interpolator, frequency, mode).
Returns a tuple ((grad_r, grad_n), HamiltonianAux(...))
where \(\partial \mathcal{H}/\partial \boldsymbol{r}\)
and \(\partial \mathcal{H}/\partial \boldsymbol{n}\) are computed in one shared
forward+backward pass, halving the number of B-interpolator evaluations compared
to computing each gradient separately.
The internally computed quantities (magnetic_field, rho,
electron_density_1e20_per_m3) are
returned as aux data, which is not differentiated.
raytrax.physics.absorption.absorption_coefficient(refractive_index: jt.Float[jax.Array, 3], magnetic_field: jt.Float[jax.Array, 3], electron_density_1e20_per_m3: ScalarFloat, electron_temperature_keV: ScalarFloat, frequency: ScalarFloat, mode: Literal['X', 'O']) -> ScalarFloat
Compute the absorption coefficient \(\alpha\).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
refractive_index
|
Float[Array, 3]
|
Refractive index vector in the lab system. |
required |
magnetic_field
|
Float[Array, 3]
|
Magnetic field vector in the lab system. |
required |
electron_density_1e20_per_m3
|
ScalarFloat
|
Electron density in 10^20 m^-3. |
required |
electron_temperature_keV
|
ScalarFloat
|
Electron temperature in keV. |
required |
frequency
|
ScalarFloat
|
Frequency of the wave in Hz. |
required |
mode
|
Literal['X', 'O']
|
Polarization mode, either "X" for extraordinary or "O" for ordinary. |
required |
Returns:
| Type | Description |
|---|---|
ScalarFloat
|
Absorption coefficient \(\alpha\) in 1/m. |
Plotting
1D
raytrax.plot.plot1d.plot_linear_power_density(beam_profile: BeamProfile, ax=None)
Plot the linear power density as function of arc length.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
beam_profile
|
BeamProfile
|
The traced BeamProfile. |
required |
ax
|
Optional matplotlib axes. If None, creates new figure. |
None
|
Returns:
| Type | Description |
|---|---|
|
Matplotlib axis. |
raytrax.plot.plot1d.plot_radial_power_density(radial_profile: RadialProfile, ax=None)
Plot the radial power deposition profile.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
radial_profile
|
RadialProfile
|
The output RadialProfile with volumetric power density. |
required |
ax
|
Optional matplotlib axis. If None, creates new figure. |
None
|
Returns:
| Type | Description |
|---|---|
|
Matplotlib axis. |
raytrax.plot.plot1d.plot_radial_electron_density(radial_profiles: RadialProfiles, ax=None)
Plot the electron density as function of rho.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
radial_profiles
|
RadialProfiles
|
The input RadialProfiles. |
required |
ax
|
Optional matplotlib axes. If None, creates new figure with two subplots. |
None
|
Returns:
| Type | Description |
|---|---|
|
Matplotlib axes (single axis or tuple of two axes). |
raytrax.plot.plot1d.plot_radial_electron_temperature(radial_profiles: RadialProfiles, ax=None)
Plot the electron temperature as function of rho.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
radial_profiles
|
RadialProfiles
|
The input RadialProfiles. |
required |
ax
|
Optional matplotlib axes. If None, creates new figure. |
None
|
Returns:
| Type | Description |
|---|---|
|
Matplotlib axis. |
2D
raytrax.plot.plot2d.plot_beamtrace_rz(beam_trace: BeamProfile, phi: float, ax=None, add_colorbar: bool = True, **kwargs)
Plot the beam trace in the R-Z plane, coloured by linear power density.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
beam_trace
|
BeamProfile
|
The traced BeamProfile. |
required |
phi
|
float
|
Toroidal angle (unused; kept for API consistency). |
required |
ax
|
Matplotlib axes. If None, a new figure is created. |
None
|
|
add_colorbar
|
bool
|
If True (default), attach a matched-height colorbar. |
True
|
**kwargs
|
Passed to LineCollection (e.g. |
{}
|
Returns:
| Type | Description |
|---|---|
|
The matplotlib axes object. |
raytrax.plot.plot2d.plot_electron_density_rz(magnetic_configuration: MagneticConfiguration, radial_profiles: RadialProfiles, phi: float, ax=None, **kwargs)
Plot the electron density in the R-Z plane.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
magnetic_configuration
|
MagneticConfiguration
|
The magnetic configuration object. |
required |
radial_profiles
|
RadialProfiles
|
Radial profiles containing electron density vs rho. |
required |
phi
|
float
|
The toroidal angle at which to evaluate. |
required |
ax
|
Matplotlib axes to plot on. If None, creates new figure. |
None
|
|
**kwargs
|
Additional arguments passed to contourf (e.g., levels, cmap). |
{}
|
Returns:
| Type | Description |
|---|---|
|
The matplotlib axes object. |
raytrax.plot.plot2d.plot_magnetic_field_rz(magnetic_configuration: MagneticConfiguration, phi: float, ax=None, **kwargs)
Plot the magnetic field strength in the R-Z plane.
raytrax.plot.plot2d.plot_effective_radius_rz(magnetic_configuration: MagneticConfiguration, phi: float, ax=None, **kwargs)
Plot the normalized effective radius rho in the R-Z plane.
3D
raytrax.plot.plot3d.plot_beam_profile_3d(beam_profile: BeamProfile, plotter=None, tube_radius=0.01, n_spline_points=100, **kwargs)
Plot a beam profile as a 3D tube colored by linear power density.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
beam_profile
|
BeamProfile
|
The beam profile to plot. |
required |
plotter
|
Optional PyVista plotter to add the mesh to. If None, creates a new one. |
None
|
|
tube_radius
|
Radius of the tube in meters (default: 0.01). |
0.01
|
|
n_spline_points
|
Number of points for spline interpolation (default: 100). |
100
|
|
**kwargs
|
Additional keyword arguments passed to plotter.add_mesh(). |
{}
|
Returns:
| Type | Description |
|---|---|
|
A PyVista plotter object. Call .show() to display. |
raytrax.plot.plot3d.plot_flux_surface_3d(magnetic_configuration: MagneticConfiguration, rho_value: float = 1.0, plotter=None, **kwargs)
Plot a 3D flux surface using PyVista.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
magnetic_configuration
|
MagneticConfiguration
|
The magnetic configuration object. |
required |
rho_value
|
float
|
The value of rho to plot (default: 1.0 for LCFS). |
1.0
|
plotter
|
Optional PyVista plotter to add the mesh to. If None, creates a new one. |
None
|
|
**kwargs
|
Additional keyword arguments passed to plotter.add_mesh(). |
{}
|
Returns:
| Type | Description |
|---|---|
|
A PyVista plotter object. Call .show() to display. |
raytrax.plot.plot3d.plot_b_surface_3d(magnetic_configuration: MagneticConfiguration, b_value: float, plotter=None, **kwargs)
Plot a 3D surface of constant magnetic field magnitude |B| using PyVista.
Useful for visualising the electron-cyclotron resonance layer, e.g. the
2nd-harmonic resonance surface at
B = f_0 / (2 × 27.99 GHz/T).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
magnetic_configuration
|
MagneticConfiguration
|
The magnetic configuration object. |
required |
b_value
|
float
|
The |B| value (T) at which to extract the isosurface. |
required |
plotter
|
Optional PyVista plotter to add the mesh to. If None, a new one is created. |
None
|
|
**kwargs
|
Additional keyword arguments forwarded to
|
{}
|
Returns:
| Type | Description |
|---|---|
|
A PyVista plotter object. Call |