Solver Settings & Performance Tuning

Solver Settings

The ODE solver used for ray tracing has five tunable parameters, exposed via TracerSettings. The defaults are chosen to work well for typical ECRH scenarios in metre-scale devices:

import raytrax

settings = raytrax.TracerSettings(
    relative_tolerance=1e-4,  # PID controller rtol
    absolute_tolerance=1e-6,  # PID controller atol
    max_step_size=0.05,  # metres
    max_arc_length=20.0,  # metres
    max_steps=4096,  # step budget, sizes the output buffers
)

result = raytrax.trace(mag_conf, profiles, beam, settings=settings)

The four float fields are ordinary JAX values (not static compile-time constants), so changing them does not trigger a recompilation. max_steps is the exception: it sizes the padded output buffers, so changing it does trigger a recompilation (see below).

When to Adjust the Settings

max_step_size is the most impactful parameter. The adaptive controller (Bosh3 + a PI PIDController) will take steps smaller than this when needed, but never larger. Reducing it forces finer integration near caustics or in steep-gradient regions — at the cost of more steps. In vacuum the error estimate is zero, so this cap is also what keeps the solver from stepping over the plasma.

relative_tolerance / absolute_tolerance control the error estimate the PID controller tries to satisfy per step. Tightening them (e.g. rtol=1e-6) improves accuracy but increases the number of steps. The defaults are sufficient for power deposition profiles; gradient-based optimisation may benefit from tighter tolerances.

max_arc_length is a safety cutoff. Rays that haven't terminated (by leaving the plasma or being absorbed) are stopped at this arc length. 20 m is generous for any current fusion device; you would only need to lower it to speed up debugging of misconfigured inputs, or raise it for very large devices.

max_steps is the step budget of the ODE solver and sizes the output buffers (max_steps + 1 entries per trajectory quantity with trim=False). If a ray needs more steps, typically because of a small max_step_size or tight tolerances, the trace stops early: TraceResult.incomplete is set and, with trim=True, trace emits a RuntimeWarning naming the cause. The absorbed power of such a trace is underestimated, so raise max_steps until the flag is no longer set. Reaching max_arc_length sets the same flag.

JAX Compilation

Raytrax is built on JAX, which compiles Python functions to optimised machine code the first time they are called. On subsequent calls with inputs of the same types and shapes, the cached compiled function is reused — making repeated calls (e.g. inside an optimisation loop) very fast.

The first call to trace is therefore noticeably slower than all subsequent ones. This is expected and is the reason Raytrax is designed to be used as a library inside a program, not as a command-line tool.

What Triggers a Recompilation

JAX recompiles whenever the structure of the inputs changes. For Raytrax, this means:

In practice, your grid shapes and device type are fixed for a given study, so you typically pay the compilation cost once per Python session.

Pre-warming the Cache

If you want to avoid the compilation delay inside a time-sensitive loop, call trace once with representative inputs before your main computation:

import raytrax

# One-time compilation warm-up
_ = raytrax.trace(mag_conf, profiles, beam)

# Subsequent calls use the cached compiled function
for beam in beam_sweep:
    result = raytrax.trace(mag_conf, profiles, beam)

Persistent Compilation Cache

Pre-warming only helps within a single Python session. JAX also supports a persistent on-disk cache that stores compiled XLA artifacts between runs, so even the very first call in a new session is fast:

import jax

jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")

import raytrax  # import after setting the cache dir

Alternatively, set the environment variable before starting Python:

export JAX_COMPILATION_CACHE_DIR=/tmp/jax_cache

The cache is keyed on the compiled computation and the JAX/XLA version, so it is invalidated automatically when you upgrade JAX. It is safe to share the cache directory across multiple scripts that use Raytrax with the same input shapes.