Optimistix

Optimistix 是一个基于 JAX 的非线性求解器库，用于求解方程的根、优化问题、不动点以及最小二乘问题。

其主要特性包括：

• 可互操作的求解器：例如，可以自动将求根问题转换为最小二乘问题，然后使用优化算法进行求解。
• 模块化的优化器：例如，可以结合 BFGS 二次模型与狗腿下降路径，并采用信赖域更新策略。
• 使用 PyTree 作为状态表示。
• 快速的编译和运行时性能。
• 与 Optax 兼容。
• 充分利用 JAX 的优势：自动微分、自动并行化、GPU/TPU 支持等。

安装

pip install optimistix

需要 Python 3.11 或更高版本。

文档

文档地址：https://docs.kidger.site/optimistix。

快速示例

import jax.numpy as jnp
import optimistix as optx

# 让我们用隐式欧拉法求解常微分方程 dy/dt = tanh(y(t))。
# 我们需要找到 y1，使得 y1 = y0 + tanh(y1) * dt。

y0 = jnp.array(1.)
dt = jnp.array(0.1)

def fn(y, args):
    return y0 + jnp.tanh(y) * dt

solver = optx.Newton(rtol=1e-5, atol=1e-5)
sol = optx.fixed_point(fn, solver, y0)
y1 = sol.value  # 满足 y1 == fn(y1)

引用

如果您在学术工作中使用了本库，请引用以下文献：(arXiv 链接)

@article{optimistix2024,
    title={Optimistix: modular optimisation in JAX and Equinox},
    author={Jason Rader and Terry Lyons and Patrick Kidger},
    journal={arXiv:2402.09983},
    year={2024},
}

相关项目：JAX 生态系统中的其他库

始终实用
Equinox：神经网络及 JAX 核心库中尚未包含的所有内容！
jaxtyping：用于数组形状和数据类型的类型注解。

深度学习
Optax：一阶梯度优化器（SGD、Adam 等）。
Orbax：检查点保存（异步/多主机/多设备）。
Levanter：可扩展且可靠的基座模型训练（如 LLM）。
paramax：PyTree 的参数化与约束。

科学计算
Diffrax：数值微分方程求解器。
Lineax：线性方程组求解器。
BlackJAX：概率与贝叶斯采样。
sympy2jax：SymPy 到 JAX 的转换；可通过梯度下降训练符号表达式。
PySR：符号回归。（非 JAX 的优秀项目！）

Awesome JAX
Awesome JAX：更全面的 JAX 项目列表。

致谢

Optimistix 主要由 Jason Rader (@packquickly) 开发：Twitter；GitHub；官网。目前由 Johanna Haffner (@johannahaffner) 协同维护：GitHub；官网。

Optimistix 快速上手指南

Optimistix 是一个基于 JAX 的非线性求解器库，支持求根、最小化、不动点迭代和最小二乘法等问题。它具备模块化设计、自动微分、GPU/TPU 加速以及与 Optax 生态互操作等特性。

环境准备

• Python 版本：3.11 或更高
• 前置依赖：
  • JAX（会自动安装）
  • 可选：Optax（用于梯度优化器集成）

💡 提示：国内用户可使用清华或中科大镜像源加速安装依赖。

安装步骤

使用 pip 直接安装：

pip install optimistix

如需使用国内镜像加速安装：

pip install optimistix -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下示例演示如何使用 optx.fixed_point 求解一个简单的不动点问题（隐式欧拉法解 ODE）：

import jax.numpy as jnp
import optimistix as optx

# 定义方程：y1 = y0 + tanh(y1) * dt
y0 = jnp.array(1.)
dt = jnp.array(0.1)

def fn(y, args):
    return y0 + jnp.tanh(y) * dt

# 选择牛顿法求解器
solver = optx.Newton(rtol=1e-5, atol=1e-5)

# 求解不动点
sol = optx.fixed_point(fn, solver, y0)
y1 = sol.value  # 满足 y1 == fn(y1)

该示例展示了 Optimistix 的核心用法：定义目标函数、选择求解器、调用求解接口并获取结果。更多高级功能（如最小化、自定义优化路径等）请参考官方文档：https://docs.kidger.site/optimistix