jaxdf - 基于JAX的离散化框架

概述 | 示例 | 安装 | 文档 | 支持

<br/>

概述

Jaxdf 是一个基于 JAX 的软件包，提供了一个编码框架，用于创建具有任意离散化的可微分数值模拟器。

Jaxdf 的主要目标是帮助构建物理系统的数值模型，如波传播或偏微分方程的数值求解，以便轻松适应用户的研究需求。这些模型是纯函数，可以无缝集成到用 JAX 编写的任意可微分程序中。例如，它们可以用作神经网络中的层，或用于构建物理损失函数。

<br/>

示例

以下脚本构建了非线性算子 (∇² + sin)，在二维正方形区域上应用傅里叶谱离散化。然后利用这个算子定义了一个损失函数。使用 JAX 的自动微分计算了这个损失函数的梯度。

[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object],[object Object]

<br/>

安装

在安装 jaxdf 之前，请确保您的系统上已经安装了 JAX。如果您打算使用支持 NVidia GPU 的 jaxdf，请按照说明相应地安装 JAX。

要从 PyPI 安装 jaxdf，请使用 pip 命令：

[object Object],[object Object],[object Object]

对于开发目的，可以通过克隆仓库或下载并解压缩压缩包来安装 jaxdf。之后，在终端中导航到根文件夹，执行以下命令：

[object Object],[object Object],[object Object],[object Object]

这将安装依赖项和软件包本身（以可编辑模式）。

支持

如果您遇到任何代码问题或希望提出新功能，请随时提出问题。如果您需要指导、想讨论某些事项，或只是想打个招呼，请随时在我们的 Discord 频道中留言。

<br/>

贡献

我们非常欢迎贡献！大多数贡献都始于一个问题。请不要犹豫，创建问题来询问功能、反馈性能或简单地联系我们。

要提交拉取请求，请查看详细的贡献指南了解如何操作，但基本上请记住以下主要准则：

• 如果您添加新功能或修复错误：
  • 确保它有测试覆盖
  • 使用 kacl-cli 在更新日志中添加一行
• 如果您更改了文档中的内容，请确保可以使用 mkdocs serve 正确构建文档站点

<br/> <br/>

引用

本软件包的初始版本在 NeurIPS 2021 的可微分编程研讨会上展示。

@article{stanziola2021jaxdf,
    author={Stanziola, Antonio and Arridge, Simon and Cox, Ben T. and Treeby, Bradley E.},
    title={A research framework for writing differentiable PDE discretizations in JAX},
    year={2021},
    journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}
}

<br/>

致谢

• 本仓库的一些打包工作是通过编辑 @rochacbruno 的这个模板完成的
• 使用的多重调度方法基于 plum，请查看这个amazing项目：https://github.com/wesselb/plum

相关项目

1. odl 算子离散化库（ODL）是一个用于快速原型设计的 Python 库，专注于（但不限于）反问题。
2. deepXDE：一个用于科学机器学习的 TensorFlow 和 PyTorch 库。
3. SciML：SciML 是一个由 NumFOCUS 赞助的开源软件组织，旨在统一科学机器学习的软件包。