jaxphys

jaxphys 是一个基于 JAX 构建的 GPU 加速可微分物理引擎。它填补了重型研究代码（如 FEniCS、OpenFOAM、COMSOL）难以安装且不可微分，以及教育库仅支持 CPU 且玩具级之间的空白。所有内容均为纯 JAX 实现，因此模拟可在 jax.jit、jax.vmap 和 jax.grad 下运行。 涵盖的领域和求解器： - 经典力学：拉格朗日系统和哈密顿系统、辛欧拉法、蛙跳法（Störmer-Verlet）、Yoshida-4、RK4、欧拉法、速度-Verlet N 体、刚体、自适应 RK45 步进器。 - 量子：一维和二维分裂算子薛定谔求解器、有限差分本征态、紧束缚能带结构、通过精确对角化的海森堡自旋链、林德blad 主方程。 - 电磁学：FDTD（二维 TM 和三维）带分裂场 PML、FDFD（二维 TM）带拉伸坐标 PML、Boris 推进电荷、矩形波导模式。 - 流体：弱可压缩 SPH、可压缩欧拉（一维 MUSCL-HLLC）、格子玻尔兹曼 D2Q9、涡量-流函数纳维-斯托克斯方程。 - 统计力学：Ising Metropolis（棋盘格）和 Wolff 簇蒙特卡洛、玻尔兹曼统计（蒙特卡洛采样本身不可微分）。 - 光学：ABCD 光线追踪、夫琅禾费衍射。 关键工作流是定义拉格朗日量，让 JAX 自动微分生成运动方程，然后通过整个轨迹进行优化。优化模块提供基于梯度的优化、敏感性分析、参数网格扫描和最佳扫描点的细化。积分器记录了阶数和辛性质；辛积分器假设可分离的哈密顿方程，因此 LagrangianSystem.simulate 仅接受 euler 和 rk4，而 HamiltonianSystem 用于辛积分。 README 报告了与向量化 NumPy 基线相比的仅 CPU 基准测试（例如 N 体 16.3 倍、FDTD 三维 8.4 倍、薛定谔二维 1.7 倍、SPH 0.6 倍，其中 SciPy cKDTree 基线在 CPU 上胜出），并注明测量来自共享的 4 核云容器，未测量 GPU 数字。验证套件检查求解器与解析或独立结果：开普勒能量误差对蛙跳和 Yoshida-4 有界，而 RK4 漂移、测量的收敛阶、从 Binder 累积量交叉得到的 Ising 临界温度、FDTD 腔模与 Yee 色散关系、薛定谔范数守恒和自由波包精度、LBM 通道衰减和泊肃叶剖面、Sod 激波管 L1 误差、FDFD 与 Hankel 格林函数、以及 jax.grad 梯度与中心有限差分比较。 安装通过 pip install jaxphys；示例涵盖双摆、弹道优化、量子隧穿、FDTD 衍射和三维偶极辐射、Ising 相变、卡门涡街、航天器轨迹瞄准和基准脚本。开发使用 pytest、ruff 和 mypy，离线演示可通过 uv 运行。项目采用 MIT 许可，并引用标准物理和数值积分参考文献。