有限基础物理信息神经网络 (FBPINN) JAX 问题

Question

我正在尝试修改 Ben Moseley 在 github 上提供的代码https://github.com/benmoseley/FBPINNs。 我的目的是将一个值向量插入到依赖于 x y 坐标的损失 fn 中，并且我需要将原始向量 Z 作为 x 和 y 的函数进行插值，然后将值与算法使用的坐标相同提取样本 x 和 y，以便值匹配。 我遇到的问题是，在 loss fn 中，我无法使用 JAX 以外的库，并且据我所知，JAX 中没有可以在 2D 中进行插值的函数。

我试图以各种方式解决这个问题，但我没有成功，我的想法之一是推断算法采样的 x,y 点，但我没有成功，代码确实非常清晰。有人能给我任何建议/帮助吗？

会有 jax.scipy.ndimage.map_coordinates 函数，但它无法正常工作，并且它推断的点毫无意义。

Answer 1

如果线性或最近邻插值就足够了，您可以使用

jax.scipy.interpolate.RegularGridInterpolator

来完成您需要的操作

如果您需要更复杂的东西，例如样条插值，则

jax

本身不包含任何内容。也就是说，您也许能够找到适合您的下游实现。我遇到的一个可能值得尝试的是在

jax_cosmo

项目中：https://jax-cosmo.readthedocs.io/en/latest/_modules/jax_cosmo/scipy/interpolate.html.