JAX-FEM4Geo

传统岩土有限元软件擅长完成正向计算，却往往难以回答另一类问题：如果监测位移发生变化，材料参数应该怎样更新？如果设计变量有成百上千个，怎样高效计算它们对目标函数的影响？JAX-FEM4Geo 将 JAX 的自动微分能力引入地质力学有限元，探索把“模拟结果”进一步变成可用于反演和优化的梯度信息。

面向岩土材料的可微塑性实现

项目基于 JAX-FEM 实现 Drucker–Prager 压力相关塑性模型，用于描述土、岩石和混凝土等材料。代码包含增量应力更新、历史应力与应变变量、塑性返回映射和有限元积分点上的向量化计算。

Drucker–Prager 屈服面在锥顶附近存在导数奇异性，这对自动微分尤其不友好。项目使用双曲正则化，将 sqrt(J₂) 替换为 sqrt(J₂ + a²)，在尽量保留物理行为的同时平滑梯度；代码还包含屈服判断与顶点返回处理，为研究“非光滑本构如何进入可微计算图”提供了具体实现。

不只给出模型，也验证计算链路

仓库提供多个层次的验证工具：

• 单材料点本构驱动程序，在无网格条件下生成单轴应力路径。
• 三维块体与圆柱体三轴压缩模拟，支持围压、逐步加载、卸载和应力不变量分析。
• FEM 结果、理论或半解析结果的应力—应变曲线对比。
• 自动微分（AD）与有限差分（FD）梯度对照，用于检查目标函数对弹性模量和黏聚参数的敏感度。
• 黏塑性、直接求解、平滑截断和硬化等多种策略实验，用于诊断塑性阶段的伴随求解稳定性。

这些内容让项目不只是一个本构类，而是一套从公式、材料点、有限元正向解到反向梯度的验证链路。研究者可以看到一个可微力学模型怎样搭建，也能看到正向收敛并不自动保证反向伴随问题稳定这一关键边界。

从正向模拟走向参数反演

可微有限元最有吸引力的地方，是反向模式自动微分可以在一次正向模拟和一次反向传播后，计算目标函数对大量参数的梯度。对于基于监测数据反演空间材料参数、校准本构模型或进行梯度优化，这可能比逐参数有限差分更具扩展潜力。

JAX-FEM4Geo 目前仍是科研代码，重点集中在 Drucker–Prager 模型、三轴验证和可微性探索，并非完整的工程分析软件。它适合从事计算岩土、可微物理、反问题或本构模型研究的人，也适合希望理解 JAX、有限元与自动微分如何衔接的开发者。项目最有价值的部分，恰恰是同时保留了成功路径和数值困难：它提供了继续研究参数反演、平滑塑性和稳定伴随算法的真实起点。