google/jax-cfd

JAX‑CFD: JAX를 이용한 미분 가능한 계산 유체 역학

개요 – JAX를 사용하여 핵심 CFD 솔버(유한체적법 및 유사 스펙트럼법)를 구현한 오픈 소스 연구용 라이브러리입니다. 따라서 코드는 자동으로 미분 가능하며 GPU/TPU에서 실행할 수 있습니다. 또한 Haiku로 구축된 몇 가지 머신러닝 증강 모델도 제공합니다.

중요성 – JAX의 JIT 컴파일과 자동 미분 기능을 활용하여 연구자들은 다음과 같은 작업을 수행할 수 있습니다:

유체 시뮬레이션과 상호작용하는 신경망을 학습시킵니다 (예: 대리 모델, 폐쇄 모델).
CFD 솔버에서 직접 경사도 기반 최적화 또는 데이터 동화를 수행합니다.
코드 변경 없이 CPU, GPU 또는 TPU에서 동일한 코드를 실행합니다.

주요 구성 요소

jax_cfd.base – 스테거드 그리드 (Arakawa‑C/MAC) 상의 유한체적법/차분법 솔버.
jax_cfd.spectral – 안티앨리어싱 기능이 포함된 유사 스펙트럼 와류 솔버.
jax_cfd.ml – ML 증강 CFD 모델 (Haiku 신경망) 및 예시 사전 학습 모델.
jax_cfd.data – 시뮬레이션 데이터 로딩, 처리 및 시각화를 위한 유틸리티 (Xarray + Pillow).

수치 계산 초점

비정상 난류 유동.
공간 이산화: 스테거드 그리드 FV/Diff 또는 유사 스펙트럼 와류.
시간 이산화: 1차 명시적 이류, 암시적/명시적 확산.
압력 솔버: Conjugate‑gradient 또는 FFT 기반 대각화 (주기적 BCs만 해당).
이류 스킴: 2차 Van Leer.
난류 폐쇄: Smagorinsky eddy‑viscosity.

시작하는 방법

핵심 패키지를 설치합니다:
```
pip install jax-cfd
```

선택 사항인 ML 또는 데이터 유틸리티를 위한 추가 기능을 추가합니다:

pip install jax-cfd[ml]   # Haiku + ML helpers
pip install jax-cfd[data] # Xarray, Pillow, etc.

notebooks/ 폴더에 있는 준비된 Colab 노트북을 실행하여 다음을 확인합니다:
- 2-D 스테거드 그리드 FV 데모
- 2-D 유사 스펙트럼 난류 데모
- 채널 플로우 예시
- PNAS 2021 논문의 결과 재현 (데이터 분석 & 모델 추론)

프로젝트 상태 – 더 이상 유지보수되지 않음 (2024년 6월)으로 표시되었습니다. README는 JAX‑Fluids, PhiFlow, 및 Exponax와 같은 활발히 유지보수되는 대체재를 가리킵니다.

전형적인 사용 사례

미분 가능한 CFD에 대한 학술 연구, 예: 난류 폐쇄 모델 학습 또는 역관측 연산자 학습.
ML‑CFD 하이브리드 모델의 프로토타이핑, 즉 유체 솔버를 통한 경사도를 필요로 하는 모델.
JAX의 유체 역학 워크로드에서의 성능 벤치마크.

인용 – 유한체적법 코드를 사용하거나 ML 모델을 사용한다면, PNAS 2021 논문을 인용하십시오. 스펙트럼법 코드를 사용한다면, arXiv 2022 pre‑print를 인용하십시오. \n관련 프로젝트 – README는 다른 미분 가능한 CFD 라이브러리 (PhiFlow, Autograd fluid examples) 및 더 넓은 JAX 과학용 도구 (JAX‑MD, JAX‑DFT, jax‑cosmo, Veros)를 목록에 나열열합니다.

위의 모든 세부 사항은 저장소의 README에서 직접 가져온 것입니다.