Working with sparse linear solvers #785
-
|
Hi all, I am building a Python framework for numerical geodynamic modelling that uses all the tricks to get high performance with simple code. One of the features that I am working on is writing a convenient python interface for a wide range of direct linear system solvers. This is extremely important as we often deal with extremely ill posed problems and it is helpful to test out different solvers quickly. All solvers I want to expose follow the Eigen sparse solver concept and have identical methods, so let's take PardisoLU for example. Solving a sparse linear system involves 3 key steps:
The models we work with require the solution of a sparse linear system at each timestep. The values of the matrix change at each iteration, but the sparsity structure is often identical over the whole simulation. This means that it is possible to perform step 1) only once and then store and reuse the computed permutation in order to gain a substantial speedup. I am now looking for the best approach to implement bindings for my Eigen solver objects. The first attempt I made was to wrap the solver object in a shared_ptr for shared ownership as well as using nb::DRef to automatically map my scipy CSC matrix to Eigen::SparseMatrix objects: #include <nanobind/nanobind.h>
#include <nanobind/stl/shared_ptr.h>
#include <nanobind/eigen/dense.h>
#include <iostream>
// Eigen modules
#include <Eigen/Sparse>
#include <Eigen/PardisoSupport>
namespace nb = nanobind;
using SparseMatrix = Eigen::SparseMatrix<double, Eigen::ColMajor>;
using SparseSolver = Eigen::PardisoLU<SparseMatrix>;
using SolverPtr = std::shared_ptr<SparseSolver>;
// Function to create a solver object, returning it as a void* for Python
SolverPtr create_solver() {
return std::make_shared<SparseSolver>();
}
// Function to analyze the pattern of the matrix
void analyze_pattern(SolverPtr solver, nb::DRef<SparseMatrix> matrix) {
solver->analyzePattern(matrix);
if (solver->info() != Eigen::Success) {
throw std::runtime_error("Symbolic factorization failed");
}
}
// Function to factorize the matrix
void factorize(SolverPtr solver, nb::DRef<SparseMatrix> matrix){
solver->factorize(matrix);
if (solver->info() != Eigen::Success) {
throw std::runtime_error("Numerical factorization failed");
}
}
// Function to solve the system
Eigen::VectorXd solve(SolverPtr solver, const nb::DRef<Eigen::VectorXd> & b) {
Eigen::VectorXd x = solver->solve(b);
if (solver->info() != Eigen::Success) {
throw std::runtime_error("Solution failed");
}
return x;
}
// Wrapper functions for Python
NB_MODULE(pardiso_support, m) {
m.def("create_solver", &create_solver,
"Create a PardisoLU solver object");
m.def("analyze_pattern", &analyze_pattern,
"Perform symbolic factorization of the matrix",
nb::arg("solver"), nb::arg("matrix"));
m.def("factorize", &factorize,
"Perform numerical factorization of the matrix",
nb::arg("solver"), nb::arg("matrix"));
m.def("solve", &solve,
"Solve the linear system",
nb::arg("solver"), nb::arg("b"));
}This however has a few issue:
To get around this, I have tried to:
However this not only feels like a hack, but also non-deterministically results in:
What is the best way to go about this? |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 3 replies
-
|
To my knowledge, |
Beta Was this translation helpful? Give feedback.
-
|
Nanobind provides support for In the shared pointer type caster ( static_assert(is_base_caster_v<Caster>,
"Conversion of ``shared_ptr<T>`` requires that ``T`` is "
"handled by nanobind's regular class binding mechanism. "
"However, a type caster was registered to intercept this "
"particular type, which is not allowed.");
|
Beta Was this translation helpful? Give feedback.
-
|
Hi Jakob, Thank you very much for the super quick reply. I understand, thanks for clarifying the issue with the sparse matrix object. Luckily, although it looks kind of ugly, I can easily get around this by passing the three internal ndarrays of the sparse matrix as Eigen::Array references and then build an Eigen::Map on the fly for each function call. The bigger problem is the SparseSolver object as it contains data that I cannot easily extract and pass around. I understand that shared pointers may not work in this case, however I am surprised that simply creating the solver object via "new" and then passing around an nb::capsule pointing to it is also problematic. In this case the C++ side will not automatically call the deconstructor of the object and from what I understand nanobind will only call it when the python wrapper is garbage collected. Yet I am getting non-deterministic segfaults and other errors. Could you elaborate on this? Many thanks in advance! |
Beta Was this translation helpful? Give feedback.
-
|
Nevermind what I said. I must've messed up something because using nb::capsule does indeed work! Except having to explicitly define #include <nanobind/nanobind.h>
#include <nanobind/eigen/dense.h>
// Eigen modules
#include <Eigen/Sparse>
#include <Eigen/PardisoSupport>
namespace nb = nanobind;
using SparseMatrix = Eigen::Map<const Eigen::SparseMatrix<double, Eigen::ColMajor>>;
using SparseSolver = Eigen::PardisoLU<SparseMatrix>;
using solver_ptr = SparseSolver *;
using matrix_ptr = SparseMatrix *;
// Function to create a solver object, returning it as a void* for Python
nb::capsule create_solver() {
solver_ptr ptr = new SparseSolver();
return nb::capsule(ptr, "PardisoSolverPtr", [](void *ptr) noexcept {
delete static_cast<solver_ptr>(ptr);
});
}
// Function to create a sparse matrix object, returning it as a void* for Python
nb::capsule create_matrix(const int N,
const nb::DRef<Eigen::VectorXi> & outer,
const nb::DRef<Eigen::VectorXi> & inner,
const nb::DRef<Eigen::VectorXd> & values) {
// Create a sparse matrix object from the input
const int nnzero = values.size();
matrix_ptr ptr = new SparseMatrix(N, // Rows
N, // Columns
nnzero, // Non-zero elements
outer.data(), // Outer indices
inner.data(), // Inner indices
values.data()); // Values
return nb::capsule(ptr, "SparseMatrixPtr", [](void *ptr) noexcept {
delete static_cast<matrix_ptr>(ptr);
});
}
// Function to analyze the pattern of the matrix
void analyze_pattern(nb::capsule solver_capsule, nb::capsule matrix_capsule) {
solver_ptr solver = static_cast<solver_ptr>(solver_capsule.data());
matrix_ptr matrix = static_cast<matrix_ptr>(matrix_capsule.data());
solver->analyzePattern(*matrix);
if (solver->info() != Eigen::Success) {
throw std::runtime_error("Symbolic factorization failed");
}
}
// Function to factorize the matrix
void factorize(nb::capsule solver_capsule, nb::capsule matrix_capsule) {
solver_ptr solver = static_cast<solver_ptr>(solver_capsule.data());
matrix_ptr matrix = static_cast<matrix_ptr>(matrix_capsule.data());
solver->factorize(*matrix);
if (solver->info() != Eigen::Success) {
throw std::runtime_error("Numerical factorization failed");
}
}
// Function to solve the system
Eigen::VectorXd solve(nb::capsule solver_capsule, const nb::DRef<Eigen::VectorXd> & b) {
solver_ptr solver = static_cast<solver_ptr>(solver_capsule.data());
Eigen::VectorXd x = solver->solve(b);
if (solver->info() != Eigen::Success) {
throw std::runtime_error("Solution failed");
}
return x;
}
// Wrapper functions for Python
NB_MODULE(pardiso_support, m) {
m.def("create_solver", &create_solver,
"Create a PardisoLU solver object");
m.def("create_matrix", &create_matrix,
"Create a sparse matrix object",
nb::arg("N"), nb::arg("outer"), nb::arg("inner"), nb::arg("values"));
m.def("analyze_pattern", &analyze_pattern,
"Perform symbolic factorization of the matrix",
nb::arg("solver"), nb::arg("matrix"));
m.def("factorize", &factorize,
"Perform numerical factorization of the matrix",
nb::arg("solver"), nb::arg("matrix"));
m.def("solve", &solve,
"Solve the linear system",
nb::arg("solver"), nb::arg("b"));
}
Thanks again! |
Beta Was this translation helpful? Give feedback.
Nevermind what I said.
I must've messed up something because using nb::capsule does indeed work!
Except having to explicitly define
noexceptwhen initialising the capsule, this approach is easy and avoids having to re-initialize objects many times. Also object ownership is very clear so this is a plus!