jaxphys

jaxphys est un moteur de physique différentiable accéléré par GPU construit sur JAX. Il cible un vide entre les codes de recherche lourds (FEniCS, OpenFOAM, COMSOL) difficiles à installer et non différentiables, et les bibliothèques éducatives limitées au CPU et de niveau jouet. Tout est en JAX pur, donc les simulations s'exécutent sous jax.jit, jax.vmap et jax.grad. Domaines et solveurs couverts : - Mécanique classique : systèmes lagrangiens et hamiltoniens, Euler symplectique, leapfrog (Störmer-Verlet), Yoshida-4, RK4, Euler, N-corps velocity-Verlet, corps rigides, pas adaptatif RK45. - Quantique : solveurs de Schrödinger à opérateur scindé en 1D et 2D, états propres par différences finies, structures de bandes en liaison forte, chaînes de spins de Heisenberg par diagonalisation exacte, équation maîtresse de Lindblad. - Électromagnétisme : FDTD (2D TM et 3D) avec PML à champ scindé, FDFD (2D TM) avec PML à coordonnées étirées, charges poussées par Boris, modes de guide d'ondes rectangulaire. - Fluides : SPH faiblement compressible, Euler compressible (1D MUSCL-HLLC), lattice Boltzmann D2Q9, Navier-Stokes en vorticité-fonction de courant. - Mécanique statistique : Ising Metropolis (damier) et Monte Carlo par amas de Wolff, statistiques de Boltzmann (l'échantillonnage Monte Carlo lui-même n'est pas différentiable). - Optique : tracé de rayons ABCD, diffraction de Fraunhofer. Un flux de travail clé consiste à définir un lagrangien et laisser l'autodiff de JAX produire les équations du mouvement, puis optimiser à travers des trajectoires entières. Le module optimize fournit une optimisation par gradient, une analyse de sensibilité, des balayages de grilles de paramètres et un raffinement des meilleurs points de balayage. Les intégrateurs sont documentés avec leur ordre et leurs propriétés symplectiques ; les intégrateurs symplectiques supposent des équations de Hamilton séparables, donc LagrangianSystem.simulate n'accepte que euler et rk4, tandis que HamiltonianSystem est utilisé pour l'intégration symplectique. Le README rapporte des benchmarks CPU uniquement contre des bases NumPy vectorisées (par exemple N-corps 16,3x, FDTD 3D 8,4x, Schrödinger 2D 1,7x, SPH 0,6x où la base SciPy cKDTree gagne sur CPU), avec la mise en garde que les mesures proviennent d'un conteneur cloud partagé à 4 cœurs et qu'aucun chiffre GPU n'a été mesuré. Une suite de validation vérifie les solveurs par rapport à des résultats analytiques ou indépendants : erreur d'énergie de Kepler bornée pour leapfrog et Yoshida-4 tandis que RK4 dérive, ordres de convergence mesurés, température critique d'Ising à partir des croisements de cumulants de Binder, modes de cavité FDTD par rapport à la relation de dispersion de Yee, conservation de la norme de Schrödinger et précision du paquet libre, décroissance de canal LBM et profil de Poiseuille, erreur L1 du tube à choc de Sod, FDFD par rapport à la fonction de Green de Hankel, et gradients jax.grad comparés aux différences finies centrées. L'installation se fait via pip install jaxphys ; les exemples couvrent un double pendule, l'optimisation de projectile, l'effet tunnel quantique, la diffraction FDTD et le rayonnement dipolaire 3D, la transition de phase d'Ising, l'allée de tourbillons de Kármán, le ciblage de trajectoire spatiale et un script de benchmark. Le développement utilise pytest, ruff et mypy, avec une démo hors ligne exécutable via uv. Le projet est sous licence MIT et cite des références standard de physique et d'intégration numérique.