返回 主页

Python调用和编译Eigen库函数

资料:

准备

  1. 安装vs2022,选择C++桌面开发;
  2. 下载Eigen源码
  3. 安装python3.10+, numpy, pybind11;

1. 编写cpp代码

在solve_KUF.cpp中:

#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <stdexcept>
namespace py = pybind11;
Eigen::VectorXd solveKUF(const Eigen::MatrixXd& A, const Eigen::VectorXd& B) {
    // 输入验证
    if (A.rows() != A.cols()) {
        throw std::invalid_argument("Matrix A must be square");
    }
    if (A.rows() != B.rows()) {
        throw std::invalid_argument("Matrix A and Vector B must have same number of rows");
    }

    // 转换为稀疏矩阵
    Eigen::SparseMatrix<double> sparseA = A.sparseView();

    // 使用BiCGSTAB求解器
    Eigen::BiCGSTAB<Eigen::SparseMatrix<double>> solver;
    solver.compute(sparseA);

    if (solver.info() != Eigen::Success) {
        throw std::runtime_error("Matrix decomposition failed");
    }

    Eigen::VectorXd X = solver.solve(B);

    if (solver.info() != Eigen::Success) {
        throw std::runtime_error("Solving failed");
    }
    return X;
}

// Pybind11模块定义
PYBIND11_MODULE(solve_KUF, m) {
    m.doc() = "Python module for solving sparse linear system AX=B using Eigen";

    m.def("solveKUF", &solveKUF,
          py::arg("A"), py::arg("B"),
          "Solve the linear system AX=B for sparse matrix A and vector B\n"
          "Args:\n"
          "    A: NxN numpy array (sparse matrix in dense format)\n"
          "    B: Nx1 numpy array\n"
          "Returns:\n"
          "    Solution vector X as numpy array");
}

pybind11支持numpy数组和eigen MatixXd之间的自动转换.

2. 编写setup.py文件

在setup.py中:

from setuptools import setup, Extension
import pybind11

# 定义扩展模块
ext_modules =Extension(
        'solve_KUF',
        sources=['solve_KUF.cpp'],
        include_dirs=[pybind11.get_include(),
                      'L:\\EF2D\\includes\\eigen-3.4.0'],  # 添加Eigen和pybind11头文件路径
        language='c++',
        extra_compile_args=['--std=c++17', '--O3']  # 优化标志

    )

setup(
    name='solve_KUF',
    version='0.1',
    ext_modules=[ext_modules],
)

3. 编译

在目录中执行命令, 如果出现build文件夹和solve_KUF.cp310-win_amd64.pyd, 说明编译成功

python.exe路径 setup.py build_ext --inplace

4. 测试

import numpy as np
import solve_KUF
# 创建一个稀疏矩阵的密集表示
N = 5
A = np.zeros((N, N))
np.fill_diagonal(A, 2.0)  # 对角线为2
A[0, 1] = 1.0
A[1, 0] = 1.0

# 创建右侧向量
B = np.random.rand(N)

# 求解AX=B
X = solve_KUF.solveKUF(A, B)

print("A:\n", A)
print("\nB:\n", B)
print("\nSolution X:\n", X)
print("\nA @ X (should equal B):\n", A @ X)

返回 主页