프로젝트 소개

jaxphys는 JAX 위에 구축된 GPU 가속 미분 가능 물리 엔진입니다. 설치가 어렵고 미분이 불가능한 대형 연구 코드(FEniCS, OpenFOAM, COMSOL)와 CPU 전용의 교육용 라이브러리 사이의 격차를 목표로 합니다. 모든 것이 순수 JAX로 작성되어 시뮬레이션이 jax.jit, jax.vmap 및 jax.grad에서 실행됩니다. 포함된 영역 및 솔버: - 고전역학: 라그랑지안 및 해밀토니안 시스템, 심플렉틱 오일러, leapfrog(Störmer-Verlet), Yoshida-4, RK4, 오일러, 속도-베를레 N체, 강체, 적응형 RK45 스텝퍼. - 양자역학: 1D 및 2D 분할 연산자 슈뢰딩거 솔버, 유한 차분 고유 상태, 타이트 바인딩 밴드 구조, 정확한 대각화를 통한 하이젠베르크 스핀 사슬, Lindblad 마스터 방정식. - 전자기학: 분할 필드 PML을 사용한 FDTD(2D TM 및 3D), 늘어난 좌표 PML을 사용한 FDFD(2D TM), Boris 추진 전하, 직사각형 도파관 모드. - 유체역학: 약한 압축성 SPH, 압축성 오일러(1D MUSCL-HLLC), 격자 볼츠만 D2Q9, 와도-유선 함수 Navier-Stokes. - 통계역학: Ising Metropolis(체커보드) 및 Wolff 클러스터 몬테카를로, 볼츠만 통계(몬테카를로 샘플링 자체는 미분 불가). - 광학: ABCD 광선 추적, 프라운호퍼 회절. 핵심 워크플로는 라그랑지안을 정의하고 JAX 자동 미분이 운동 방정식을 생성한 다음 전체 궤적을 통해 최적화하는 것입니다. optimize 모듈은 그래디언트 기반 최적화, 민감도 분석, 매개변수 그리드 스윕 및 최상의 스윕 포인트 정제를 제공합니다. 적분기는 차수와 심플렉틱 속성으로 문서화되어 있습니다. 심플렉틱 적분기는 분리 가능한 해밀턴 방정식을 가정하므로 LagrangianSystem.simulate는 euler와 rk4만 허용하는 반면 HamiltonianSystem은 심플렉틱 적분에 사용됩니다. README는 벡터화된 NumPy 기준선에 대한 CPU 전용 벤치마크를 보고합니다(예: N체 16.3배, FDTD 3D 8.4배, 슈뢰딩거 2D 1.7배, SPH 0.6배, 여기서 SciPy cKDTree 기준선이 CPU에서 우세). 측정이 공유 4코어 클라우드 컨테이너에서 이루어졌고 GPU 수치가 측정되지 않았다는 주의 사항이 있습니다. 검증 스위트는 솔버를 해석적 또는 독립적 결과와 대조합니다: leapfrog 및 Yoshida-4에 대한 케플러 에너지 오차 경계(RK4는 드리프트), 측정된 수렴 차수, Binder 누적량 교차에서 Ising 임계 온도, Yee 분산 관계에 대한 FDTD 공동 모드, 슈뢰딩거 노름 보존 및 자유 패킷 정확도, LBM 채널 감쇠 및 Poiseuille 프로파일, Sod 충격 튜브 L1 오차, Hankel 그린 함수에 대한 FDFD, 그리고 중앙 유한 차분과 비교한 jax.grad 그래디언트. 설치는 pip install jaxphys를 통해 이루어집니다. 예제는 이중 진자, 발사체 최적화, 양자 터널링, FDTD 회절 및 3D 쌍극자 복사, Ising 상전이, Kármán 와류 거리, 우주선 궤적 목표 지정 및 벤치마크 스크립트를 다룹니다. 개발은 pytest, ruff 및 mypy를 사용하며 uv를 통해 실행 가능한 오프라인 데모가 있습니다. 프로젝트는 MIT 라이선스이며 표준 물리 및 수치 적분 참조를 인용합니다.