JAX-FEM｜当有限元遇上机器学习

近年来，机理与数据的融合计算正快速发展着。机器学习作为数据驱动的计算方法，以其解决高维复杂问题的能力而闻名于世。它在视觉理解、自然语言处理等领域取得了重大突破。与此同时，基于机理的计算方法，如有限元模拟，在现代工业中起着至关重要的作用，广泛应用于飞机、汽车、材料和新能源等行业。

有限元方法简介与挑战

有限元法（FEM，Finite Element Method）是一种为求解偏微分方程边值问题近似解的数值技术。求解时对整个问题区域进行分解，每个子区域上定义基函数，把有限元组合起来近似得到原方程的解。然而，被广泛应用的有限元方法也面临挑战：随着问题规模的增长，计算复杂度和内存需求也随之增加，导致计算时间过长和资源消耗过大（“维度灾难”问题）。

逆问题是在数学和科学中，通过给定输出数据，寻找导致这些结果的输入或参数。求解逆问题的价值在于提供无法直接观测的信息和改善决策预测。有限元在许多逆问题中是常用的求解方法，然而逆问题中，用有限元法求解会有很多难点：
1. 非唯一性：可能存在多个不同的输入或参数组合对应相同的观测结果，无法得到唯一解决方案。

2. 不适定问题：输入数据中的误差会放大，导致求解结果不稳定。
3. 计算复杂度高：特别是在高维度或复杂模型下，计算复杂度显著增加。

---

机器学习在 FEM 中的应用与痛点

机器学习（Machine Learning, ML）通过从数据中总结规律，并使用大量参数进行学习，有潜力更好地处理复杂、大规模数据，并在非线性建模上有优势，推动科学和工程进步，高效地解决"维数灾难"。例如，神经网络和深度学习可以用于预测、优化和模拟 FEM 计算过程，以及解决逆问题。然而，在实际应用中，ML 算法面临一些挑战。由于其"黑箱"性质、训练困难、分布外难泛化和数据缺乏等问题，机器学习方法无法直接应用于严肃科学场景。于是机理与数据融合计算就变得尤为重要，它试图在保持物理准确性和性能之间取得平衡，使人们对科学的认识可以转化为大规模可应用的程序与系统。

            数据与机理融合计算，图片来自Physics-informed machine learning

然而，机理与数据融合计算在实际应用中遇到了许多困难，新兴的机器学习框架和传统有限元框架往往不兼容。初学者甚至有经验的专家们常常会遇到以下问题：

1. Python 程序速度过慢而 C/C++/Fortran 开发与调试难度过高；
2. 手动计算微分和 Sensitivity 非常麻烦且容易犯错；
3. 在研究数据驱动的计算问题时，需要使用多个工具并在它们之间频繁传输数据，非常繁琐。

---

Notebook 上手实践

它来了它来了，JAX-FEM: A differentiable GPU-accelerated 3D finite element solver 带着它的三大优势来了：

1. Life Is Short, Use Python! JAX-FEM 使用纯 Python 实现，同时可利用 GPU 以高效地解决中等到大规模问题。

2. 利用 JAX 的可微编程实现高效便捷的逆向设计。
   

                                逆问题：拓扑优化示例

3. 可以与机器学习无缝集成。无须在不同框架下切换开发。
   

                    JAX-FEM实现的数据驱动的多尺度计算方案工作流示意图

   现在点击下方图片来开始阅读吧！

---

欢迎大家来 Notebook 案例广场，获取更多有意思的实践～