PySR 寻找能够优化特定目标的符号表达式。

https://github.com/MilesCranmer/PySR/assets/7593028/c8511a49-b408-488f-8f18-b1749078268f

PySR：Python 和 Julia 中的高性能符号回归

文档	论坛	论文	Colab 示例

pip	conda	统计信息
		pip: conda:

如果您觉得 PySR 很有用，请引用论文 arXiv:2305.01582。 如果您已经使用 PySR 完成了某个项目，请提交 PR，在 研究展示页面 上分享您的成果！

目录：

• 为什么选择 PySR？
• 安装
• 快速入门
• → 文档
• 贡献者

测试状态

Linux	Windows	macOS
Docker	Conda	覆盖率

为什么选择 PySR？

PySR 是一款用于 符号回归 的开源工具：这是一种机器学习任务，其目标是找到一个可解释的符号表达式，以优化某个目标函数。

在数年的时间里，PySR 从零开始设计，旨在做到（1）尽可能高效，（2）尽可能灵活可配置，以及（3）易于使用。PySR 与 Julia 库 SymbolicRegression.jl 联合开发，后者构成了 PySR 强大的搜索引擎。这些算法的详细内容已在 PySR 论文 中阐述。

符号回归最适合低维数据集，但也可以通过“神经网络符号蒸馏”将其扩展到高维空间，如 2006.11287 所述，我们将其应用于 N 体问题。本质上，这种方法利用符号回归将神经网络转换为解析方程。因此，这些工具同时提供了一种明确而强大的方式来解释深度神经网络。

安装

pip

您可以通过 pip 安装 PySR：

pip install pysr

Julia 的依赖项将在首次导入时自动安装。

conda

同样地，使用 conda：

conda install -c conda-forge pysr

docker build -t pysr .

docker run -it --rm pysr ipython

apptainer build --notest pysr.sif Apptainer.def

apptainer run pysr.sif

export LD_LIBRARY_PATH=$HOME/.julia/juliaup/julia-1.10.0+0.x64.linux.gnu/lib/julia/:$LD_LIBRARY_PATH

快速入门

您可以尝试这里的交互式教程 这里，它使用了 examples/pysr_demo.ipynb 中的笔记本。

在实际操作中，我强烈建议使用 IPython 而不是 Jupyter，因为它的输出效果更好。下面是一个快速演示，您可以将其粘贴到 Python 运行环境中。

首先，我们导入 numpy 来生成一些测试数据：

import numpy as np

X = 2 * np.random.randn(100, 5)
y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5

我们创建了一个包含 100 个样本、每个样本有 5 个特征的数据集。我们希望建模的关系是 $2.5382 \cos(x_3) + x_0^2 - 0.5$。

现在，让我们创建一个 PySR 模型并进行训练。PySR 的主要接口风格类似于 scikit-learn：

from pysr import PySRRegressor

model = PySRRegressor(
    maxsize=20,
    niterations=40,  # < 增加此值以获得更好的结果
    binary_operators=["+", "*"],
    unary_operators=[
        "cos",
        "exp",
        "sin",
        "inv(x) = 1/x",
        # ^ 自定义运算符（Julia 语法）
    ],
    extra_sympy_mappings={"inv": lambda x: 1 / x},
    # ^ 同时为 SymPy 定义运算符
    elementwise_loss="loss(prediction, target) = (prediction - target)^2",
    # ^ 自定义损失函数（Julia 语法）
)

这将设置模型进行 40 次迭代的搜索，每次迭代包含数十万次突变和方程评估。

让我们用我们的数据集来训练这个模型：

model.fit(X, y)

在内部，这会启动一个 Julia 进程，该进程将以多线程方式搜索适合数据集的方程。

训练过程中会打印出方程，当您满意时，可以通过按 'q' 然后按回车键提前退出。

模型拟合完成后，您可以运行 model.predict(X) 来查看使用自动选择的表达式对给定数据集的预测结果；或者例如，运行 model.predict(X, 3) 来查看第 3 个方程的预测结果。

您还可以运行以下命令来打印学习到的方程：

print(model)

输出如下：

PySRRegressor.equations_ = [
	   pick     score                                           equation       loss  complexity
	0        0.000000                                          4.4324794  42.354317           1
	1        1.255691                                          (x0 * x0)   3.437307           3
	2        0.011629                          ((x0 * x0) + -0.28087974)   3.358285           5
	3        0.897855                              ((x0 * x0) + cos(x3))   1.368308           6
	4        0.857018                ((x0 * x0) + (cos(x3) * 2.4566472))   0.246483           8
	5  >>>>       inf  (((cos(x3) + -0.19699033) * 2.5382123) + (x0 *...   0.000000          10
]

pick 列中的箭头表示当前由您的 model_selection 策略选定用于预测的方程。（您也可以在 .fit(X, y) 之后更改 model_selection。）

model.equations_ 是一个 pandas DataFrame，包含了所有方程，包括可调用格式 (lambda_format)、SymPy 格式 (sympy_format - 您也可以通过 model.sympy() 获取）、甚至 JAX 和 PyTorch 格式（两者都可微分 - 您可以通过 model.jax() 和 model.pytorch() 获取）。

请注意，PySRRegressor 会保存上次搜索的状态，并在您下次调用 .fit() 时从上次停止的地方继续，前提是您已设置 warm_start=True。如果对搜索参数进行了重大更改（例如更改运算符），则可能会出现问题。您可以运行 model.reset() 来重置状态。

您会注意到 PySR 会保存两个文件：hall_of_fame...csv 和 hall_of_fame...pkl。CSV 文件列出了方程及其损失，而 PKL 文件则是模型的保存状态。您可以从 PKL 文件加载模型，方法如下：

model = PySRRegressor.from_file("hall_of_fame.2022-08-10_100832.281.pkl")

此外，PySR 还提供了其他一些有用的功能，例如去噪（如 denoise=True）和特征选择（如 select_k_features=3）。有关这些功能及其他功能的示例，请参阅 示例页面。有关更多选项的详细信息，请参阅 选项页面。您还可以在 API 页面 上查看完整的 API。此外，调优页面 提供了调整 PySR 的技巧。

详细示例

以下代码尽可能地使用了 PySR 的所有功能。请注意，这只是一个功能演示，您不应直接使用此示例。有关每个参数的具体作用，请参阅 API 页面。

model = PySRRegressor(
    populations=8,
    # ^ 假设我们有4个核心，这意味着每个核心运行2个种群，因此总有一个种群在持续运行。
    population_size=50,
    # ^ 稍大的种群规模，以增加多样性。
    ncycles_per_iteration=500,
    # ^ 种群之间迁移的代数间隔。
    niterations=10000000,  # 无限运行
    early_stop_condition=(
        "stop_if(loss, complexity) = loss < 1e-6 && complexity < 10"
        # 如果找到一个既好又简单的方程，则提前停止
    ),
    timeout_in_seconds=60 * 60 * 24,
    # ^ 或者，在24小时后停止。
    maxsize=50,
    # ^ 允许更高的复杂度。
    maxdepth=10,
    # ^ 但避免过深的嵌套。
    binary_operators=["*", "+", "-", "/"],
    unary_operators=["square", "cube", "exp", "cos2(x)=cos(x)^2"],
    constraints={
        "/": (-1, 9),
        "square": 9,
        "cube": 9,
        "exp": 9,
    },
    # ^ 限制每个参数内的复杂度。
    # "inv": (-1, 9) 表示分子没有约束，
    # 但分母的最大复杂度为9。
    # "exp": 9 则表示 `exp` 只能接受复杂度为9的表达式作为输入。
    nested_constraints={
        "square": {"square": 1, "cube": 1, "exp": 0},
        "cube": {"square": 1, "cube": 1, "exp": 0},
        "exp": {"square": 1, "cube": 1, "exp": 0},
    },
    # ^ 对运算符进行嵌套约束。例如，
    # "square(exp(x))" 是不允许的，因为 "square": {"exp": 0}。
    complexity_of_operators={"/": 2, "exp": 3},
    # ^ 自定义特定运算符的复杂度。
    complexity_of_constants=2,
    # ^ 惩罚常数比变量更重。
    select_k_features=4,
    # ^ 只使用最重要的4个特征进行训练。
    progress=True,
    # ^ 如果输出到文件，可以设置为False。
    weight_randomize=0.1,
    # ^ 更频繁地随机化树结构。
    cluster_manager=None,
    # ^ 可以设置为例如 "slurm" 来运行 slurm 集群。只需从主节点启动一个脚本即可。
    precision=64,
    # ^ 使用更高精度的计算。
    warm_start=True,
    # ^ 从上次中断的地方继续。
    turbo=True,
    # ^ 更快的评估（实验性）
    extra_sympy_mappings={"cos2": lambda x: sympy.cos(x)**2},
    # extra_torch_mappings={sympy.cos: torch.cos},
    # ^ 这里不需要，因为 cos 已经定义，但这正是定义自定义 PyTorch 运算符的方式。
    # extra_jax_mappings={sympy.cos: "jnp.cos"},
    # ^ 对于 JAX，需要传递字符串。
)

Docker

您也可以在 Docker 中测试 PySR，而无需在本地安装它。只需在本仓库的根目录下运行以下命令：

docker build -t pysr .

这将为您的系统架构构建一个名为 pysr 的镜像，其中还包含 IPython。您可以选择特定版本的 Python 和 Julia，方法如下：

docker build -t pysr --build-arg JLVERSION=1.10.0 --build-arg PYVERSION=3.11.6 .

然后，您可以使用以下 Dockerfile 运行容器：

docker run -it --rm -v "$PWD:/data" pysr ipython

这会将当前目录挂载到容器的 /data 目录，并启动 IPython。

如果您在构建适用于您系统架构的镜像时遇到问题，可以在 build 和 run 命令前添加 --platform linux/amd64 来模拟其他架构。

我们热忱欢迎新贡献者！请查看我们的贡献者 指南，获取实用建议 🚀。如果您对新功能有任何想法，请随时在 issues 或 discussions 页面上分享。

Mark Kittisopikul 💻 💡 🚇 📦 📣 👀 🔧 ⚠️	T Coxon 🐛 💻 🔌 💡 🚇 🚧 👀 🔧 ⚠️ 📓	Dhananjay Ashok 💻 🌍 💡 🚧 ⚠️	Johan Blåbäck 🐛 💻 💡 🚧 📣 👀 ⚠️ 📓	JuliusMartensen 🐛 💻 📖 🔌 💡 🚇 🚧 📦 📣 👀 🔧 📓	ngam 💻 🚇 📦 👀 🔧 ⚠️	Christopher Rowley 💻 💡 🚇 📦 👀	Kaze Wong 🐛 💻 💡 🚇 🚧 📣 👀 🔬 📓
Christopher Rackauckas 🐛 💻 🔌 💡 🚇 📣 👀 🔬 🔧 ⚠️ 📓	Patrick Kidger 🐛 💻 📖 🔌 💡 🚧 📣 👀 🔬 🔧 ⚠️ 📓	Okon Samuel 🐛 💻 📖 🚧 💡 🚇 👀 ⚠️ 📓	William Booth-Clibborn 💻 🌍 📖 📓 🚧 👀 🔧 ⚠️	Pablo Lemos 🐛 💡 📣 👀 🔬 📓	Jerry Ling 🐛 💻 📖 🌍 💡 📣 👀 📓	Charles Fox 🐛 💻 💡 🚧 📣 👀 🔬 📓	Johann Brehmer 💻 📖 💡 📣 👀 🔬 ⚠️ 📓
Marius Millea 💻 💡 📣 👀 📓	Coba 🐛 💻 💡 👀 📓	foxtran 💻 💡 🚧 🔧 📓	Shah Mahdi Hasan 🐛 💻 👀 📓	Pietro Monticone 🐛 📖 💡	Mateusz Kubica 📖 💡	Jay Wadekar 🐛 💡 📣 🔬	Anthony Blaom, PhD 🚇 💡 👀
Jgmedina95 🐛 💡 👀	Michael Abbott 💻 💡 👀 🔧	Oscar Smith 💻 💡	Eric Hanson 💡 📣 📓	Henrique Becker 💻 💡 👀	qwertyjl 🐛 📖 💡 📓	Rik Huijzer 💡 🚇	Hongyu Wang 💡 📣 🔬
Zehao Jin 🔬 📣	Tanner Mengel 🔬 📣	Arthur Grundner 🔬 📣	sjwetzel 🔬 📣 📓	Saurav Maheshkar 🔧

PySR 快速上手指南

PySR 是一个高性能的符号回归（Symbolic Regression）工具，旨在从数据中发现可解释的数学表达式。它基于 Julia 的强大搜索引擎，同时提供 Python 风格的接口（类似 scikit-learn）。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 建议 Python 3.8+
• 依赖说明:
  • PySR 依赖 numpy, pandas, sympy 等常见 Python 库。
  • Julia 环境: 首次导入 PySR 时，它会自动下载并配置所需的 Julia 版本及 SymbolicRegression.jl 包，无需手动安装 Julia。
• 硬件建议: 支持多线程并行搜索，多核 CPU 能显著提升发现公式的速度。

安装步骤

你可以选择 pip 或 conda 进行安装。国内用户若遇到下载慢的问题，可配置相应的镜像源。

方式一：使用 pip 安装

pip install pysr

国内加速建议:

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

方式二：使用 conda 安装

conda install -c conda-forge pysr

国内加速建议 (使用清华/中科大镜像):

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
conda config --set show_channel_urls yes
conda install pysr

注意: 安装完成后，第一次在 Python 中 import pysr 时会自动触发 Julia 环境的初始化，请耐心等待下载完成。

基本使用

PySR 的使用流程与 scikit-learn 高度相似：定义模型 -> 拟合数据 -> 获取结果。

1. 准备数据与定义模型

以下示例生成一组测试数据，并配置 PySR 搜索包含加减乘除、三角函数和自定义算子的表达式。

import numpy as np
from pysr import PySRRegressor

# 1. 生成测试数据
# 真实关系为: y = 2.5382 * cos(x_3) + x_0^2 - 0.5
X = 2 * np.random.randn(100, 5)
y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5

# 2. 定义模型
model = PySRRegressor(
    maxsize=20,          # 表达式最大复杂度
    niterations=40,      # 迭代次数 (实际应用中建议增大此值以获得更好结果)
    binary_operators=["+", "*"],
    unary_operators=[
        "cos",
        "exp",
        "sin",
        "inv(x) = 1/x",  # 自定义算子 (Julia 语法)
    ],
    extra_sympy_mappings={"inv": lambda x: 1 / x}, # 为 SymPy 定义对应的 Python 映射
    elementwise_loss="loss(prediction, target) = (prediction - target)^2", # 自定义损失函数
)

2. 训练模型

调用 fit 方法启动搜索过程。底层会启动 Julia 进程进行多线程搜索。 在训练过程中，终端会实时打印发现的方程。如需提前停止，可在终端按下 q 然后回车。

model.fit(X, y)

3. 查看结果与预测

训练完成后，可以查看发现的公式列表或直接进行预测。

# 打印学习到的方程列表
print(model)

# 使用自动选择的最佳方程进行预测
predictions = model.predict(X)

# 或者指定使用第 3 个发现的方程进行预测
predictions_3 = model.predict(X, 3)

输出示例解读: print(model) 将输出一个表格，其中 pick 列带有 >>>> 标记的行表示当前被选中用于预测的最佳方程。表格还包含了方程的复杂度（complexity）和损失值（loss）。

PySRRegressor.equations_ = [
	   pick     score                                           equation       loss  complexity
	0        0.000000                                          4.4324794  42.354317           1
	1        1.255691                                          (x0 * x0)   3.437307           3
	...
	5  >>>>       inf  (((cos(x3) + -0.19699033) * 2.5382123) + (x0 *...   0.000000          10
]

进阶提示

• 断点续训: PySR 会自动保存状态文件（.pkl 和 .csv）。下次运行 fit 时，若设置 warm_start=True（默认行为），将从上次中断处继续搜索。
• 格式导出: 结果可导出为 SymPy (model.sympy())、JAX (model.jax()) 或 PyTorch (model.pytorch()) 格式，便于后续微分或部署。
• 加载模型: 可通过 PySRRegressor.from_file("path/to/file.pkl") 直接加载已训练好的模型。