3D Visualisation

import pyvista as pv
import jax.numpy as jnp

from raytrax import Beam, RadialProfiles, trace
from raytrax.examples.w7x import (
    PortA,
    get_w7x_magnetic_configuration,
    w7x_aiming_angles_to_direction,
)
from raytrax.plot.plot3d import plot_flux_surface_3d, plot_beam_profile_3d

mag_conf = get_w7x_magnetic_configuration()

rho = jnp.linspace(0, 1, 200)
profiles = RadialProfiles(
    rho=rho,
    electron_density=0.5 * (1.0 - rho**2),
    electron_temperature=3.0 * (1.0 - rho**2),
)
beam = Beam(
    position=jnp.array(PortA.D1.cartesian),
    direction=jnp.array(
        w7x_aiming_angles_to_direction(-10.0, 0.0, PortA.D1.phi_deg)
    ),
    frequency=jnp.array(140e9),
    mode="O",
    power=1e6,
)
result = trace(mag_conf, profiles, beam)

plotter = pv.Plotter()
plot_flux_surface_3d(mag_conf, rho_value=1.0, plotter=plotter, opacity=0.25)
plot_beam_profile_3d(result.beam_profile, plotter=plotter, tube_radius=0.02)

plotter.export_html("scene.html")  # save interactive HTML