jaxphys

jaxphys é um motor de física diferenciável acelerado por GPU construído sobre JAX. Ele visa preencher uma lacuna entre códigos de pesquisa pesados (FEniCS, OpenFOAM, COMSOL) difíceis de instalar e não diferenciáveis, e bibliotecas educacionais que são apenas para CPU e de nível brinquedo. Tudo é JAX puro, então as simulações rodam sob jax.jit, jax.vmap e jax.grad. Domínios e solucionadores cobertos: - Mecânica clássica: sistemas Lagrangianos e Hamiltonianos, Euler simplético, leapfrog (Störmer-Verlet), Yoshida-4, RK4, Euler, N-corpos velocity-Verlet, corpos rígidos, integrador adaptativo RK45. - Quântica: solucionadores de Schrödinger de operador dividido em 1D e 2D, autoestados por diferenças finitas, estruturas de banda tight-binding, cadeias de spin de Heisenberg via diagonalização exata, equação mestra de Lindblad. - Eletromagnetismo: FDTD (2D TM e 3D) com PML de campo dividido, FDFD (2D TM) com PML de coordenadas esticadas, cargas impulsionadas por Boris, modos de guia de onda retangular. - Fluidos: SPH fracamente compressível, Euler compressível (1D MUSCL-HLLC), lattice Boltzmann D2Q9, Navier-Stokes vorticidade-função de corrente. - Mecânica estatística: Ising Metropolis (tabuleiro de xadrez) e Monte Carlo de cluster de Wolff, estatística de Boltzmann (a amostragem de Monte Carlo em si não é diferenciável). - Óptica: traçado de raios ABCD, difração de Fraunhofer. Um fluxo de trabalho chave é definir um Lagrangiano e deixar o autodiff do JAX produzir as equações de movimento, depois otimizar através de trajetórias inteiras. O módulo optimize fornece otimização baseada em gradiente, análise de sensibilidade, varreduras de grade de parâmetros e refinamento dos melhores pontos da varredura. Os integradores são documentados com ordem e propriedades simpléticas; integradores simpléticos assumem equações de Hamilton separáveis, então LagrangianSystem.simulate aceita apenas euler e rk4, enquanto HamiltonianSystem é usado para integração simplética. O README relata benchmarks apenas de CPU contra baselines vetorizados de NumPy (por exemplo, N-corpos 16,3x, FDTD 3D 8,4x, Schrödinger 2D 1,7x, SPH 0,6x onde o baseline SciPy cKDTree vence na CPU), com a ressalva de que as medições vêm de um contêiner de nuvem compartilhado de 4 núcleos e nenhum número de GPU foi medido. Uma suíte de validação verifica os solucionadores contra resultados analíticos ou independentes: erro de energia de Kepler limitado para leapfrog e Yoshida-4 enquanto RK4 deriva, ordens de convergência medidas, temperatura crítica de Ising a partir de cruzamentos de cumulante de Binder, modos de cavidade FDTD contra a relação de dispersão de Yee, conservação da norma de Schrödinger e precisão de pacote livre, decaimento de canal LBM e perfil de Poiseuille, erro L1 do tubo de choque de Sod, FDFD versus a função de Green de Hankel, e gradientes jax.grad comparados com diferenças finitas centrais. A instalação é via pip install jaxphys; os exemplos cobrem um pêndulo duplo, otimização de projétil, tunelamento quântico, difração FDTD e radiação de dipolo 3D, transição de fase de Ising, rua de vórtices de Kármán, mira de trajetória de espaçonave e um script de benchmark. O desenvolvimento usa pytest, ruff e mypy, com uma demonstração offline executável através do uv. O projeto é licenciado sob MIT e cita referências padrão de física e integração numérica.