기술 요약: TokaGrad – 엔드 투 엔드 미분 가능한 토카막 시뮬레이터

1. 문제 정의

자기 융합 연구가 과학적 타당성 단계를 넘어 원자로 규모의 상업화로 전환됨에 따라, 새로운 원자로 개념의 설계와 고성능 플라즈마 시나리오의 최적화가 매우 중요해졌습니다. 기존의 시나리오 개발 방식은 비용이 많이 드는 시행착오, 전문가 주도의 파형 튜닝, 그리고 저차원 파라미터 스캔에 크게 의존하고 있습니다. 통합 모델링 도구(예: ASTRA, TRANSP)가 가상 실험 능력을 향상시켰음에도 불구하고, 시나리오 설계는 여전히 "인간이 개입하는(human-in-the-loop)" 과정에 머물러 있습니다.

현재 방법론의 주요 한계점은 다음과 같습니다:

• 블랙박스 최적화: 많은 머신러닝 접근 방식(베이지안 최적화, 강화 학습 등)은 시뮬레이터를 블랙박스로 취급합니다. 이는 반복적인 순방향 평가와 간접적인 민감도 추론을 요구하며, 고차원, 다중 액추에이터, 전체 방전(full-discharge) 시나리오의 경우 계산 비용이 지나치게 높아지는 결과를 초래합니다.
• 미분 불가능성: 전통적인 수치 시뮬레이션은 단계별 절차로 실행되므로, 전체 파이프라인을 통해 그래디언트(gradient)를 전파할 수 없습니다. 이는 내부 민감도를 이용한 직접적인 그래디언트 기반 제어 및 최적화를 방해합니다.
• 불완전한 전체 시나리오 모델링: 기존의 미분 가능한 수송 솔버(예: TORAX)는 주로 코어 수송과 규정된 시나리오에 집중해 왔습니다. 이들은 L-H 전이, 에지 페데스탈(edge pedestal) 형성, 그리고 진화하는 플라즈마 형상과 평형 재구성 간의 결합을 포함하는 동적인 전체 방전 시나리오를 자기 일관성 있게 모델링하는 데 한계가 있습니다.

2. 방법론

저자들은 JAX 프레임워크를 기반으로 구축된 엔드 투 엔드 미분 가능한 토카막 수송 시뮬레이터인 TokaGrad를 소개합니다. 이 시스템은 초기 조건부터 최종 플라즈마 성능에 이르는 전체 시뮬레이션 파이프라인을 하나의 연결된 미분 가능한 계산 그래프로 표현합니다. 이를 통해 자동 미분(AD)이 전체 시뮬레이션 전반에 걸쳐 자코비안(Jacobian)을 계산함으로써, 기계 파라미터, 액추에이터 파형, 그리고 플라즈마 응답을 연결할 수 있습니다.

핵심 아키텍처

TokaGrad는 다음의 미분 가능한 모듈들을 단일 그래프 내에 통합합니다:

1. 수송 프레임워크 (Transport Framework): 핵심은 정규화된 반경 방향 플럭스 좌표($\rho$)에서의 전자/이온 온도($T_{e,i}$), 전자 밀도($n_e$), 폴로이달 자기 플럭스($\psi$)에 대한 결합된 1차원 수송 방정식입니다.

   • 수치 기법: 안정성을 위해 준암시적 유한 체적법(semi-implicit finite-volume method)을 사용하며, JAX의 JIT 컴파일과 호환되는 미분 가능한 삼중 대각 행렬 솔버(Thomas 알고м)를 통해 해결됩니다.
   • 방정식: 입자 수송, 보존적 에너지 수송(for $T_{e,i}$), 그리고 부트스트랩 및 전류 구동 소스에 결합된 유도 플럭스 진화를 포함합니다.

2. 기하학 및 평형 (Geometry and Equilibrium):

   • EQDSK 모드: 현실적인 벤치마킹을 위해 외부에서 생성된 평형(예: EFIT)을 지원하지만, 기하학적 구조에 대한 그래디언트는 파라미터화되지 않는 한 제한적입니다.
   • Miller/Moment 모드: 해석적 플럭스 표면 파라미터화(대반경, 신장도, 삼각형도, 샤프라노프 이동)를 사용하는 축소된 고정 경계 모델입니다. 이 모드는 완전히 미분 가능하여, 기계 형상 파라미터($R_0, a, \kappa$)의 최적화를 가능하게 합니다.

3. 수송 폐쇄 (Transport Closures):

   • 난류 수송 (Turbulent Transport): 계산 비용이 높은 자이로키네틱 솔버를 대체하기 위해 신경망(NN) 대리 모델(TGLF-NN, QLKNN)을 사용합니다. 이러한 대리 모델은 플라즈마 프로파일로부터 수송 계수로의 미분 가능한 매핑을 제공합니다.
   • 뉴클레오클래식 수송 (Neoclassical Transport): 해석적 공식(Sauter-Angioni)과 NN 대리 모델(NEO-NN)을 모두 지원합니다.
   • 경험적 모델: 빠른 베이스라인 설정을 위해 Bohm-gyroBohm (BgB) 모델을 포함합니다.

4. 소스 및 싱크 (Sources and Sinks):

   • 가열 (Heating): 오믹 가열, 알파 가열(Bosch-Hale 피팅 사용), 충돌 교환 및 복사 손실을 포함합니다. 보조 가열은 매끄러운 정규화된 가우시안 프로파일을 통해 모델링됩니다.
   • 입자 소스: 비미분적인 하드 스위칭을 피하기 위해 밀도 피드백 제어(Greenwald 분율 추적) 또는 규정된 형상 모델을 구현합니다.

5. L-H 전이 및 페데스탈 (L-H Transition and Pedestal):

   • 전이: 분리면(separatrix)을 통과하는 전력($P_{sep}$)을 경험적 임계값($P_{LH}$, 예: Martin 또는 Delabie 스케일링)과 비교하여 L-모드와 H-모드 사이의 전이를 제어하는 매끄러운 시그모이드 함수를 사용합니다. 이는 비미분적인 이진 스위칭을 방지합니다.
   • 페데스탈: 알파 임계 축소 모델 또는 EPED1-NN 대리 모델을 사용하여 페데스탈 타겟을 구성합니다. 페데스탈 프로파일은 smoothmax 함수를 사용하여 코어 수송 솔루션과 매끄럽게 블렌딩되며, 이를 통해 그래디언트가 전이 과정을 통해 전파되도록 보장합니다.

핵심 기여

• 최초의 자기 일관적인 미분 가능한 전체 방전 시뮬레이터: TokaGrad는 액추에이터, 플라즈마 평형, 페데스탈 형성, 그리고 구속 영역 전이(L-to-H)가 미분 가능한 그래프 내에서 함께 진화하는 동적인 전체 방전 시나리오를 모델링할 수 있는 최초의 프레임워크입니다.
• 엔드 투 엔드 미분 가능성: 고충실도 물리 폐쇄 모델(TGLF, EPED1, NEO)을 JAX 호환 NN 대리 모델로 변환함으로써, 최종 성능 지표로부터 수송, 난류, 평형 모델을 거쳐 설계 및 제어 입력까지 그래디언트가 전파될 수 있도록 합니다.
• 대화형 실시간 시뮬레이션: 이 프레임워크는 JIT 컴파일 후 롱 펄스 시뮬레이션에서 TORAX 대비 10배 이상의 속도 향상을 달는 것을 입증하였으며, 이를 통해 인과 관계의 대화형 탐색이 가능합니다.
• 그래디언트 기반 최적화: 유한 차분 스캔 없이도 액추에이터 파형, 시나리오 궤적, 원자로 설계 파라미터의 직접적인 최적화를 용이하게 합니다.

3. 결과 및 시연

논문은 벤치마크와 최적화 사례 연구를 통해 TokaGrad를 검증합니다:

• 벤치마킹:

  • TORAX 대비: 원형 기하 구조를 가진 단순화된 L-모드 시나리오에서, TokaGrad는 확립된 JAX 기반 TORAX 시뮬레이터와 약 5% 오차 범위 내에서 코어 수송 프로파일(온도, 밀도, 전류)을 재현했습니다.
  • 성능: Toka-Grad는 JIT 컴파일 오버헤드가 상쇄된 후, 롱 펄스 시뮬레이션에서 TORAX보다 10배 이상의 월 시간(wall-clock time) 속도 향상을 보여주었습니다.
  • H-모드 시뮬레이션: TGLF-NN/EPED1-NN 및 QLKNN/alpha-critical 모델을 사용한 ITER 베이스라인 시나리오 시뮬레이션은 약 7초 만에(실시간보다 빠른 속도로) 핵융합 이득 $Q \sim 11$ 및 핵융합 출력 $\sim 550$ MW를 달성했습니다.

• 전체 방전 L-to-H 시뮬레이션:

  • 150초간의 ITER 베이스라인 시뮬레이션은 램프업(ramp-up), L-모드 단계, 그리고 H-모드로의 전이를 성공적으로 모델링했습니다.
  • 시뮬레이션은 L-H 전이의 히스테리시스(hysteresis)를 포착했습니다: 과도한 가열 오버슈트(73 MW)가 전이를 유발했으며, 이후 보조 가열을 50 MW로 줄였음에도 알파 가열이 H-모드를 유지시켰습니다.

• 최적화 응용:

  • 액추에이터 제어: 그래디언트 기반 제어(Adam 옵티마이저)를 통해 안정성 제약 조건($q_{95} > 3, \beta_N < 3$)을 준수하면서 핵융합 지표($Q + 0.01 P_{fus}$)를 극대화하도록 플라즈마 전류($I_p$), 보조 가열($P_{aux}$), Greenwald 분율($f_{GW}$)을 성공적으로 조정했습니다.
  • 시나리오 파형 최적화: H-모드에 도달하지 못하는 초기 나이브한 램프업 상태에서 시작하여, 옵티마이저는 자동으로 과도한 가열 오버슈트 전략을 발견했습니다. 이를 통해 동일한 최종 액추에이터 조건에서도 플라즈마가 L-H 임계값을 통과하여 고성능 H-모드를 유지할 수 있게 했습니다.
  • 원자로 설계 최적화: 프레임워크는 핵융합 성능을 극대화하기 위해 기계 설계 파라미터($R_0, a, B_T, I_p$)를 최적화했습니다. 옵티마이저는 대반경 및 자기장의 증가와 같은 경향을 정확히 식별하는 동시에, 안정성 제약 조건을 만족하도록 파라미터를 조정했습니다.

4. 의의 및 주장

저자들은 TokaGrad가 자동화된 토카막 시나리오 설계를 향한 실질적인 경로를 제공한다고 주장합니다. 미분 가능한 물리학을 통해 가능해진 그래디언트 기반 방법을 블랙박스 최적화 대신 사용함으로써, 이 프레임워크는 다음을 가능하게 합니다:

• 직접적인 민감도 분석: 외부 섭동 없이 가열 궤적이 핵융합 이득에 어떻게 의존하는지, 또는 기하학적 구조가 페데스탈 높이에 어떻게 영향을 미치는지 직접 계산할 수 있습니다.
• 경로 의존적 최적화: 단순히 최종 상태뿐만 아니라, H-모드와 같은 고성능 영역에 도달하는 데 필요한 전체 동적 경로(예: 램프업 파형)를 최적화할 수 있습니다.
• 확장성: 노트북급 하드웨어(Apple M1 Pro)에서 반복적인 순방향 및 역방향 평가를 수행할 수 있는 능력은, 이전에는 계산적으로 불가능했던 복잡한 다목적 원자로 시나리오를 최적화하는 것을 현실적으로 만들어 줍니다.

결론적으로, 현재의 시연은 ITER 베이스라인과 유사한 조건에 집중되어 있으나, 이 프레임워크는 범용적이며 연소 플라즈마 시나리오 및 차세대 원자로 개념의 자동화된 물리 기반 최적화를 위한 새로운 길을 제시합니다.