tridiax

tridiax实现了在jax中的三对角系统求解器。所有求解器都支持CPU和GPU，与jit编译兼容，并且可以使用grad进行微分。

实现的求解器

• Thomas算法
• 分而治之
• Stone算法

通常，Thomas算法在CPU上会更快，而分而治之算法和Stone算法在GPU上会更快。

已知限制

目前，所有求解器仅针对维度是2的指数的系统进行了测试。

使用方法

from tridiax import thomas_solve, divide_conquer_solve, stone_solve

dim = 1024
diag = jnp.asarray(np.random.randn(dim))
upper = jnp.asarray(np.random.randn(dim - 1))
lower = jnp.asarray(np.random.randn(dim - 1))
solve = jnp.asarray(np.random.randn(dim))
solution = thomas_solve(lower, diag, upper, solve)

如果许多相同大小的系统需要求解且使用分而治之算法，则预先计算重排索引会很有帮助

from tridiax import divide_conquer_solve, divide_conquer_index

dim = 1024
diag = jnp.asarray(np.random.randn(dim))
upper = jnp.asarray(np.random.randn(dim - 1))
lower = jnp.asarray(np.random.randn(dim - 1))
solve = jnp.asarray(np.random.randn(dim))

indexing = divide_conquer_index(dim)
solution = divide_conquer_solve(lower, diag, upper, solve, indexing=indexing)

安装

git clone https://github.com/mackelab/tridiax.git
cd tridiax
pip install -e .