JAX 기반의 미분 가능한 전산 유체 역학(CFD) 프레임워크 AeroJAX

실용적 조언

역설계 프로젝트 시 AeroJAX와 같은 미분 가능한 솔버를 사용하면 유동 결과에서 형상 파라미터로 직접 그래디언트를 전달할 수 있다.
솔버 내부에 뉴럴 오퍼레이터를 통합하여 물리 법칙을 준수하는 하이브리드 모델을 구축할 때 활용 가능하다.

섹션별 상세

AeroJAX는 전체 솔버 스택을 JAX 네이티브로 구현하여 최적화 및 학습 파이프라인 내에서 엔드투엔드 미분 가능성을 유지한다. 속도, 압력, 와도 필드 전체에 걸쳐 그래디언트가 전파되도록 설계되어 물리 구조를 유지하면서도 딥러닝 모델과 결합이 가능하다. 외부 종속성 없는 CPU 우선 벡터화 구현을 통해 하드웨어 유연성을 확보했다.

2D 비압축성 나비에-스토크스 솔버와 격자 볼츠만 방법(LBM D2Q9)을 모두 지원하여 다양한 유체 해석 시나리오에 대응한다. 투영법(Projection method)과 압력 보정 기법을 사용하여 수치적 안정성을 확보했으며, Brinkman 방식의 포싱과 부드러운 마스크를 통해 복잡한 기하학적 형상을 미분 가능한 방식으로 처리한다. 이는 형상 변경에 따른 유동 변화를 수학적으로 추적할 수 있게 한다.

성능 측정 결과 128x128 격자에서 약 560 FPS, 512x96 격자에서 약 300 FPS의 실행 속도를 기록했다. 현재는 CPU 연산에 최적화되어 있으며 격자 크기에 따라 성능이 결정되는 구조이다. 솔버 루프 내부에 뉴럴 오퍼레이터나 학습된 보정 항을 삽입할 수 있는 훅(Hooks)을 제공하여 하이브리드 물리 모델 구축을 지원한다.

기존 CFD 파이프라인이 솔버를 블랙박스로 취급하여 그래디언트 기반 설계가 어려웠던 문제를 해결하는 데 중점을 두었다. 기하학적 형상에서 유동으로 이어지는 매핑 과정의 그래디언트를 역전파하여 역설계를 수행하거나, 난류 모델을 솔버 내부에서 직접 학습시키는 것이 가능하다. 유체 역학 시스템을 위한 미분 가능한 데이터 생성기로서의 가치가 확인됐다.

용어 해설

미분 가능 프로그래밍(Differentiable Programming): — 프로그램 전체의 연산 과정을 미분 가능한 형태로 구성하여 입력값에 대한 출력값의 기울기를 계산할 수 있게 하는 기법이다. 이를 통해 물리 시뮬레이터 자체를 신경망 학습 루프에 포함시켜 최적화나 역설계에 활용할 수 있다.
나비에-스토크스 방정식(Navier-Stokes Equations): — 점성을 가진 유체의 운동을 설명하는 비선형 편미분 방정식이다. 유체 역학의 핵심 방정식으로, AI 모델이 물리 법칙을 준수하며 유동을 예측하거나 제어하기 위한 학습 데이터 생성 및 물리 기반 학습의 기초가 된다.
격자 볼츠만 방법(Lattice Boltzmann Method): — 유동 현상을 거시적인 방정식 대신 미시적인 입자의 충돌과 이동으로 모델링하는 수치 해석 기법이다. 병렬화가 용이하여 GPU 가속에 적합하며, 복잡한 경계 조건을 처리하는 데 강점이 있어 최근 ML 기반 유체 시뮬레이션에 자주 활용된다.
역설계(Inverse Design): — 원하는 성능 목표를 먼저 설정하고 이를 달성하기 위한 최적의 기하학적 형상이나 물성을 찾아내는 설계 방식이다. 미분 가능한 솔버를 사용하면 성능 오차의 그래디언트를 형상 파라미터까지 역전파하여 효율적인 최적화가 가능하다.

언급된 도구

JAX추천

고성능 수치 연산 및 자동 미분 라이브러리