[논문 리뷰] TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

J. Citrin, Ian Goodfellow외 12명|arXiv (Cornell University)|2024. 06. 10.

Magnetic confinement fusion research인용 수 6

한 줄 요약

TORAX는 JAX로 구축된 Python 기반의 오픈 소스이면서 미분 가능한 1D 토카마 핵심 수송 시뮬레이터로, 빠른 실행과 기울기 기반 최적화를 가능하게 하며 RAPTOR와 비교 벤치마크된다.

ABSTRACT

We present TORAX, a new, open-source, differentiable tokamak core transport simulator implemented in Python using the JAX framework. TORAX solves the coupled equations for ion heat transport, electron heat transport, particle transport, and current diffusion, incorporating modular physics-based and ML models. JAX's just-in-time compilation ensures fast runtimes, while its automatic differentiation capability enables gradient-based optimization workflows and simplifies the use of Jacobian-based PDE solvers. Coupling to ML-surrogates of physics models is greatly facilitated by JAX's intrinsic support for neural network development and inference. TORAX is verified against the established RAPTOR code, demonstrating agreement in simulated plasma profiles. TORAX provides a powerful and versatile tool for accelerating research in tokamak scenario modeling, pulse design, and control.

연구 동기 및 목표

• Python/JAX로 구현된 미분 가능하고 빠른 핵심 토카마 수송 시뮬레이터를 제공한다.
• ML 대리모듈 및 기울기 기반 최적화 워크플로우와의 결합을 가능하게 한다.
• 정확성과 신뢰성을 확인하기 위해 TORAX를 확립된 코드(예: RAPTOR)와 벤치마크한다.
• 시나리오 모델링과 제어를 지원하기 위한 기하학, 수송, 소스, 준고전 물리학 등의 모듈식 물리 모델을 제공한다.
• 더 많은 물리학과 동적 평형을 포함하도록 확장 로드맵을 제시한다.

제안 방법

• 정규화 플럭스 좌표에서 이온/전자 열, 전자 밀도 및 전류 확산에 대한 결합된 1D 수송 편미분방정식(PDE) 세트를 형식화한다.
• 균일한 1D 격자에서 유한체적 방법으로 이산화하고 면 값에 멱법(Péclet) 가중치를 적용한다.
• theta-방법 기반 시간 적분을 사용하여 시간 진화를 해결하고, JAX를 통한 자동 미분으로 선형, 뉴턴-랩슨, 최적화 해를 제공한다.
• 기하학(CHEASE 또는 Circular), 수송(Constant, CGM, QLKNN10D), 준고전 물리(Sauter 모델) 등 모듈식 물리 모델을 도입한다.
• JAX 추론 코드와 해석자 안정성을 위한 스무딩 옵션을 갖춘 ML 대체(QLKNN10D) 결합을 가능하게 한다.

Figure 1: Poloidal cross section of the core region of a toroidal tokamak plasma, with closed nested magnetic flux surfaces. R is the major radius from the center of the torus, Z is the vertical direction. ITER-like magnitudes are shown. The flux surfaces are elliptical in the poloidal plane for ill

실험 결과

연구 질문

• RQ1 TORAX가 기울기 기반 워크플로우에 적합한 미분 가능하고 빠른 해를 사용해 핵심 토카마 수송 물리를 재현할 수 있는가?
• RQ2 RAPTOR와 같은 확립된 수송 코드와 플라즈마 프로필에 대해 TORAX의 일치성은 얼마나 좋은가?
• RQ3 모듈식 물리 모델(기하학, 수송, 준고전)들이 시뮬레이션 속도, 안정성, 확장성에 미치는 영향은 무엇인가?
• RQ4 미분 가능성과 수치적 강건성을 유지하면서 TORAX가 수송 모델에 대한 ML 대리모듈을 얼마나 통합할 수 있는가?
• RQ5 TORAX의 물리 범위와 시나리오를 확장하기 위한 로드맵 확장 계획은 무엇인가?

주요 결과

• TORAX는 JAX를 사용하여 Python에서 미분 가능하고 빠른 핵심 수송 시뮬레이터를 제공하여 기울기 기반 최적화와 ML 결합을 가능하게 한다.
• TORAX는 RAPTOR에 대해 검증되었으며 시뮬레이션된 플라즈마 프로필에서 합의를 보여준다.
• TORAX는 모듈식 프레임워크 내에서 여러 물리 모듈(기하학, 수송 모델 CGM/QLKNN10D, 준고전 Sauter)을 지원한다.
• 해법은 선형, 뉴턴-랩슨, 최적화 기반 옵션을 제공하며 로버스트를 위한 적응 타임스텝과 선 탐색/백트래킹을 갖춘다.
• JAX는 자동 미분 및 백엔드 유연성(CPU/GPU)을 가능하게 하여 ML 대리모듈 및 다중 워크플로우 통합을 촉진한다.
• 코드베이스는 오픈 소스이며 추가 물리학과 동적 평형을 포함하도록 확장 계획이 있다.

Figure 2: 1D spatial discretization with the finite volume method, with $N=5$ and $\hat{\rho}=0.2$ . The cell volume centered at $\hat{\rho}=0.5$ is highlighted. Keeping with 3D FVM formalism, a cell volume is defined with $V=d\hat{\rho}A$ , with arbitrary face area $A$ . The corresponding cell face