TORAX is a differentiable tokamak core transport simulator written in Python-JAX, designed for fast and accurate forward modeling, pulse-design, trajectory optimization, and controller design workflows. It offers open-source and extensible features, with auto-differentiation capabilities and fast runtime compilation through JAX. This allows for gradient-based nonlinear PDE solvers, trajectory optimization, and data-driven parameter identification. TORAX includes physics features such as coupled PDEs, various solvers, boundary conditions, fusion power, and more. It is not an officially supported Google product. Short-term goals include time-dependent geometry and forward sensitivity calculations. Future features may include sawtooth modeling and neoclassical tearing modes. Installation and running instructions are provided, along with plans for simulation tutorials and citing TORAX in future publications.
https://github.com/google-deepmind/torax