Dymos中使用Jax自动微分的偏导数错误排查求助

模型背景 #

我在Dymos中构建了动力学模型，使用Jax通过自动微分计算偏导数。核心代码通过自定义OpenMDAO ExplicitComponent实现，利用jax.vmap和jax.jacfwd批量计算原函数与偏导数。

偏导数检查异常现象 #

运行prob.check_partials(method='cs', compact_print=True)时，发现偏导数存在显著且系统性的误差；但使用前向差分法时误差则小很多。此外还有两个异常：

• 在_compute_primal函数中，输入参数k和I无法正常传递，运行check_partials时会变为0，必须手动赋值才能避免除零错误；
• 原模型与最小工作示例中出现误差的偏导数项不一致。

额外疑问 #

不清楚phase中的rhs_disc和rhs_col组件有何区别，希望得到解答。

求助需求 #

附上完整的模型代码（包括最小工作示例与完整船舶动力学模型），请帮忙分析问题原因，是否是模型本身或Jax使用方式导致的问题？

内容的提问来源于stack exchange，提问作者LordCat