The Thomson Scattering diagnostic is used to infer plasma conditions in fusion-relevant experiments at LLE, ZAP Energy, among others. It is a model-based diagnostic where the forward model is capable of describing the measured signal relatively accurately. TSADAR is a package that uses automatic differentiation to compute the gradients of the forward model [1]. In this example we define a Tesseract that wraps the TSADAR package by exposing its primal evaluation and the partial derivates with the apply and the jacobain endpoint respectively.
As a specific showcase we perform a gradient based parameter optimization. Intially, we pick a set of random parameters and infer the electron spectrum using the TSADAR model. Then we “rediscover” those parameters by using the L-BFGS-B algorithm implemented in scipy.optimize.minimize
to fit the parameters of the forward model to the measured signal. In the below gif we show the trajectory of the electron spectrum as the optimization progresses on the left side. On the right side we show the loss function as a function of the optimization step.
The source code can be downloaded here. tsadar.zip (3.3 MB)
Citation
- Milder, A. L., Joglekar, A. S., Rozmus, W. & Froula, D. H. Qualitative and quantitative enhancement of parameter estimation for model-based diagnostics using automatic differentiation with an application to inertial fusion. Mach. Learn.: Sci. Technol. 5, 015026 (2024).