jaxphys

jaxphys es un motor de física diferenciable acelerado por GPU construido sobre JAX. Apunta a una brecha entre códigos de investigación pesados (FEniCS, OpenFOAM, COMSOL) que son difíciles de instalar y no son diferenciables, y bibliotecas educativas que son solo para CPU y de nivel juguete. Todo es JAX puro, por lo que las simulaciones se ejecutan bajo jax.jit, jax.vmap y jax.grad. Dominios y solucionadores cubiertos: - Mecánica clásica: sistemas lagrangianos y hamiltonianos, Euler simpléctico, leapfrog (Störmer-Verlet), Yoshida-4, RK4, Euler, N-cuerpos velocity-Verlet, cuerpos rígidos, integrador adaptativo RK45. - Cuántica: solucionadores de Schrödinger de operador dividido en 1D y 2D, autoestados por diferencias finitas, estructuras de bandas de tight-binding, cadenas de espín de Heisenberg mediante diagonalización exacta, ecuación maestra de Lindblad. - Electromagnetismo: FDTD (2D TM y 3D) con PML de campo dividido, FDFD (2D TM) con PML de coordenadas estiradas, cargas impulsadas por Boris, modos de guía de onda rectangular. - Fluidos: SPH débilmente compresible, Euler compresible (1D MUSCL-HLLC), lattice Boltzmann D2Q9, Navier-Stokes vorticidad-función de corriente. - Mecánica estadística: Ising Metropolis (tablero de ajedrez) y Monte Carlo de clúster de Wolff, estadística de Boltzmann (el muestreo Monte Carlo en sí no es diferenciable). - Óptica: trazado de rayos ABCD, difracción de Fraunhofer. Un flujo de trabajo clave es definir un lagrangiano y dejar que la diferenciación automática de JAX produzca las ecuaciones de movimiento, luego optimizar a través de trayectorias completas. El módulo optimize proporciona optimización basada en gradientes, análisis de sensibilidad, barridos de cuadrícula de parámetros y refinamiento de los mejores puntos del barrido. Los integradores están documentados con orden y propiedades simplécticas; los integradores simplécticos asumen ecuaciones de Hamilton separables, por lo que LagrangianSystem.simulate acepta solo euler y rk4, mientras que HamiltonianSystem se usa para integración simpléctica. El README reporta benchmarks solo de CPU contra líneas base de NumPy vectorizado (por ejemplo, N-cuerpos 16.3x, FDTD 3D 8.4x, Schrödinger 2D 1.7x, SPH 0.6x donde la línea base SciPy cKDTree gana en CPU), con la advertencia de que las mediciones provienen de un contenedor en la nube compartido de 4 núcleos y no se midieron números de GPU. Una suite de validación verifica los solucionadores contra resultados analíticos o independientes: error de energía de Kepler acotado para leapfrog y Yoshida-4 mientras RK4 deriva, órdenes de convergencia medidos, temperatura crítica de Ising a partir de cruces de cumulantes de Binder, modos de cavidad FDTD contra la relación de dispersión de Yee, conservación de la norma de Schrödinger y precisión de paquete libre, decaimiento de canal LBM y perfil de Poiseuille, error L1 del tubo de choque de Sod, FDFD versus la función de Green de Hankel, y gradientes de jax.grad comparados con diferencias finitas centrales. La instalación es mediante pip install jaxphys; los ejemplos cubren un péndulo doble, optimización de proyectiles, tunelamiento cuántico, difracción FDTD y radiación dipolar 3D, transición de fase de Ising, calle de vórtices de Kármán, apuntado de trayectoria de naves espaciales y un script de benchmark. El desarrollo usa pytest, ruff y mypy, con una demostración sin conexión ejecutable a través de uv. El proyecto tiene licencia MIT y cita referencias estándar de física e integración numérica.