Conditional Generative Surrogate Modeling of Particle Trajectories in Stellarator Configurations

Fast evaluation of particle confinement across large stellarator design spaces is limited by the cost of repeatedly simulating particle trajectories. We investigate conditional generative modeling as a surrogate for these simulations. Particle trajectories in Boozer coordinates, together with the parallel velocity, are represented using complex short-time Fourier transforms (STFTs), and a conditional flow-matching model is trained directly in this time-frequency representation. The model is conditioned jointly on the particle initial state and stellarator configuration, and uses a structured stochastic source constructed from random-walk signals centered around the initial state. Generated spectrograms are transformed back to the time domain, allowing confinement quantities to be computed from complete generated trajectories rather than predicted directly. We evaluate the model on 40 stellarator configurations excluded from training. While particle-level escape time prediction remains challenging, with a Spearman correlation of $0.626$, aggregation over particles yields substantially stronger agreement: stellarator-level mean escape time and confinement rate achieve Spearman correlations of $0.817$ and $0.823$ respectively. These results indicate that conditional trajectory generation can capture useful configuration-dependent confinement structure and provide a promising surrogate for aggregate stellarator evaluation.