jaxdf - 基于JAX的离散化框架

概述

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

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

示例

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

from jaxdf import operators as jops
from jaxdf import FourierSeries, operator
from jaxdf.geometry import Domain
from jax import numpy as jnp
from jax import jit, grad


# 定义算子
@operator
def custom_op(u, *, params=None):
  grad_u = jops.gradient(u)
  diag_jacobian = jops.diag_jacobian(grad_u)
  laplacian = jops.sum_over_dims(diag_jacobian)
  sin_u = jops.compose(u)(jnp.sin)
  return laplacian + sin_u

# 定义离散化
domain = Domain((128, 128), (1., 1.))
parameters = jnp.ones((128,128,1))
u = FourierSeries(parameters, domain)

# 定义可微分损失函数
@jit
def loss(u):
  v = custom_op(u)
  return jnp.mean(jnp.abs(v.on_grid)**2)

gradient = grad(loss)(u) # gradient 是一个 FourierSeries

安装

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

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

pip install jaxdf

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

pip install --upgrade poetry
poetry install

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

支持

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

贡献

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

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

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

引用

本软件包的初始版本在 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}
}