yancc: A GPU-accelerated, differentiable solver for neoclassical transport in tokamaks and stellarators

기술 요약: yancc: 토카막 및 스텔라레이터를 위한 GPU 가속, 미분 가능한 뉴클리어 네오클래식 수송 솔버

문제 정의
뉴클래식 수송 이론은 자기 구속 핵융합 장치(토카막 및 스텔라레이터) 내에서 입자, 열, 운동량의 충돌 수송을 설명한다. 드리프트 키네틱 방정식(DKE)을 푸는 것은 radial electric field, bootstrap current, 그리고 수송 계수를 결정하는 데 필수적이다. 그러나 DKE의 수치적 해법은 강한 대류 지배성(advection dominance), 순환 흐름, 내부 경계층(특히 포획된 입자와 통과하는 입자 집단 사이), 심각한 이방성, 그리고 고차원성(4차원: 2개의 공간 차원, 2개의 속도 차원)으로 인해 매우 어렵다.

기존의 솔버들은 물리적 충실도와 속도 사이의 절충안에 직면해 있다. 단일 에너지(monoenergetic) 코드(예: DKES, MONKES)는 성능을 위해 종 간 결합(interspecies coupling)과 에너지 산란을 희생한다. 전체 4D 솔버(예: SFINCS, NEO)는 높은 정확도를 제공하지만, 불규칙한 메모리 접근을 동반하는 희소 행렬 분해(sparse matrix factorization)로 인해 메모리 사용량이 매우 크고 GPU 하드웨어로의 확장성이 떨어진다. 또한, 기존 코드들의 자동 미분 기능 부재는 스텔라레이터 설계 및 불확실성 정량화에 점점 더 요구되는 경사 기반 최적화 및 어드조인트 민감도 분석(adjoint sensitivity analysis)을 저해한다.

방법론
저자들은 GPU 가속과 자동 미분을 활용하기 위해 JAX 프레임워크를 기반으로 구축된 새로운 솔버인 yancc를 제시한다. 이 코드는 전체 4D DKE(속도 의존적 충돌, 에너지 산란, 완전한 종 간 결합 유지)와 축소된 단일 에너지 형태를 모두 해결한다.

• 이산화(Discretization):

  • 속도 ($x$): $e^{-x^2}$에 대해 직교하는 가중 맥스웰 다항식(weighted Maxwell polynomials)을 사용하는 콜로케이션 그리드(collocation grid)를 사용하여, 단 5~10개의 점만을 필요로 한다.
  • 피치 각도 ($\alpha$), 폴로이달 각도 ($\theta$), 토로이달 각도 ($\zeta$): 유한 차분법(finite difference)으로 이산화한다.
  • 스텐실 수정(Stencil Modification): 멀티그리드 효율성을 높이기 위해, 저자들은 표준 업윈드(upwind) 스텐실 대신 수정된 업윈드 스텐실(예: 2차 정확도를 위해 $\{0, 1, 4\}$를 사용하고, 4차 정확도를 위해 $\{-2, 0, 1, 3, 4\}$를 사용)을 채택한다. 이는 주요 차수 오차 상수를 크게 증가시키지 않으면서 결과적인 선형 시스템 행렬의 대각 우세성(diagonal dominance)을 높인다.
  • 충돌 연산자(Collision Operator): 장-입자 산란(Rosenbluth potentials)은 그린 함수 접근법을 통해 해석적으로 처리된다. 이는 레전드르 전개(Legendre expansions)와 불완전 감마 함수를 사용하여 속도 공간에서의 관련 포아송 방정식을 해결함으로써, 포텐셜을 추가적인 미지수로 풀 필요를 없앤다.

• 선형 솔버(Linear Solver):

  • 이산화된 DKE는 거대한 희소 선형 시스템 $Af = s$를 형성한다.
  • 희소 직접 솔버(sparse direct solvers) 대신, yancc는 기하학적 멀티그리드 스킴에 의해 프리컨디셔닝된 Krylov 부공간 방법(내부 GMRES를 포함한 GCROT)을 사용한다.
  • 멀티그리드 전략:
    • 그리드 계층: 세미 코서닝(semi-coarsening)이 사용된다. 속도 좌표($x$)는 모든 레벨에서 가장 미세한 해상도를 유지하는 반면, $\alpha, \theta, \zeta$는 2배로 코서닝된다.
    • 스무딩(Smoothing): 연산자의 블록 대각 근사(block-diagonal approximation)를 사용하며, 한 번에 하나의 좌표에 대해 완전한 결합을 유지하는 순열을 순환한다. 수정된 업윈드 스텐실은 이 블록 대각 스무더의 효과를 향상시킨다.
    • 조립 격자 보정(Coarse Grid Correction): 낮은 메모리 사용량을 유지하고 매트릭스 프리(matrix-free) 조립 격자 연산의 높은 비용을 피하기 위해 직접 재이산화(direct re-discretization)를 선호한다.
    • 사이클: 전체 4D DKE의 경우 V-사이클(index 1)을 사용하며, 단일 에너지 케이스의 경우 W-사이클 또는 그 이상의 사이클(index 3)이 효율적인 것으로 나타났다.

주요 기여

1. GPU 가속 전체 4D 솔버: yancc는 SFINCS와 같은 희소 행렬 기반 코드의 메모리 및 확장성 제한을 극복하고, 전체 4D DKE(완전한 종 간 결합 및 에너지 산란 포함) 솔루션과 효율적인 GPU 실행을 결합한 최초의 솔버이다.
2. 미분 가능성: JAX로 구축된 이 솔버는 완전히 미분 가능하여, 수동 미분 없이도 자기 기하 구조의 경사 기반 최적화 및 어드조인트 민감도 분석을 가능하게 한다.
3. 알고리즘 혁신: 대각 우세성을 높이기 위한 확장된 업윈드 스텐실의 도입과 DKE의 특유한 이방성에 맞춘 세미 코서닝 및 블록 대각 스무더의 사용.
4. 해석적 Rosenbluth 포텐셜: 사전 계산된 그린 함수와 텐서 수축(tensor contractions)을 사용하여 효율적인 장-입자 충돌 연산자를 구현함으로써, 그리드 상에서 보조 포아송 방정식을 푸는 계산 오버헤드를 제거하였다.

결과 및 벤치마크

• 정확도: MONKES(단일 에너지) 및 SFINCS(전체 4D)와의 벤치마크 결과, 광범위한 충돌도($\nu^* \in [10^{-4}, 10^2]$), 자기 기하 구조(W7-X, NCSX 유사 모델), 그리고 플라즈마 구성(단일 및 다종)에 대해 1% 이내의 일치도를 보였다.
• SFINCS 대비 성능:
  • 속도: yancc는 단일 스캔 기준으로 SFINCS보다 약 10배 빠른 속도를 달성한다. SFINCS의 실행 시간은 높은 충돌도에서 프리컨디셔너 성능 저하로 인해 크게 증가하는 반면, yancc의 실행 시간은 거의 일정하게 유지된다.
  • 메모리: yancc는 테스트된 시나리오에서 SFINCS(>50 GB)에 비해 약 10배 적은 메모리(예: ~6 GB)를 사용하여 GPU 기반 워크플로우에 적합하다.
• MONKES 대비 성능: 단일 에너지 벤치마크에서, yan와cc는 대부분의 충돌도 범위에서 JAX로 재구현된 MONKES 알고리즘보다 2~4배 빠르지만, 낮은 해상도가 필요한 매우 단순한 기하 구조에서는 MONKES가 더 빠를 수 있다.

의의 및 주장
본 논문은 속도, 낮은 메모리 사용량, 그리고 미분 가능성의 결합이 yancc를 현대적인 스텔라레이터 최적화 워크플로, 불확실성 정량화(몬테카를로 샘플링을 통한), 그리고 통합 프로파일 예측에 독보적으로 적합하게 만든다고 주장한다. 사전 계산된 데이터베이스나 피팅 함수에 의존하는 기존 방식과 달리, yancc는 수송 솔버 내에서 온더플라이(on-the-fly)로 뉴클래식 계수를 계산할 수 있다. 입력 파라미터에 대한 수송 계수의 경사를 계산할 수 있는 능력은 전체 구성 공간에 대한 효율적인 최적화를 가능하게 한다. 저자들은 yancc를 고충실도 물리와 설계 및 분석의 계산적 요구 사이의 간극을 메우는 도구로 위치시키며, 향후 작업으로 자기 일관적(self-consistent) radial electric field 계산 및 적응형 메쉬 세분화(adaptive mesh refinement)를 포함할 예정이라고 밝혔다.