added linear solver

Author: teseoch

README.md

@@ -1,6 +1,6 @@
 # polysolve-python
 
-Small Python binding for PolySolve's nonlinear solver interface.
+Small Python binding for PolySolve's nonlinear and linear solver interfaces.
 
 ```python
 import numpy as np
@@ -20,7 +20,7 @@ class Quadratic(polysolve.Problem):
         return 2.0 * scipy.sparse.eye(x.size, format="csc")
 
 
-x, result = polysolve.minimize(
+x, log = polysolve.minimize(
     Quadratic(),
     np.zeros(3),
     {
@@ -32,7 +32,29 @@ x, result = polysolve.minimize(
 )
 
 print(x)
-print(result)
+print(log)
 ```
 
 Python subclasses must implement `value(x)`, `gradient(x)`, and `hessian(x)`. Optional PolySolve callbacks such as `solution_changed`,  `stop`, `is_step_valid`, and `max_step_size` can also be implemented on the subclass.
+
+## Linear solves
+
+For a one-off linear system:
+
+```python
+A = scipy.sparse.csc_matrix([[4.0, 1.0], [1.0, 3.0]])
+b = np.array([1.0, 2.0])
+
+x = polysolve.solve(A, b, {"solver": "Eigen::SimplicialLDLT"})
+```
+
+For repeated solves with the same matrix pattern:
+
+```python
+solver = polysolve.LinearSolver({"solver": "Eigen::SimplicialLDLT"})
+solver.analyze_pattern(A)
+solver.factorise(A)  # factorize(A) is also available
+
+x = solver.solve(b)
+print(solver.info())
+```

cmake/polysolve.cmake

@@ -8,4 +8,4 @@ endif()
 message(STATUS "Third-party: creating target 'polysolve'")
 
 include(CPM)
-CPMAddPackage("gh:polyfem/polysolve#550f6a51fd1cf2904ac76b748c6bd0d3c20aed89")
+CPMAddPackage("gh:polyfem/polysolve#59cd6d6271d73a69236e7c714bfc580527b33fac")

example.py

@@ -17,7 +17,7 @@ def hessian(self, x):
     def post_step(self, iter_num, solver_info, x, grad):
         print(f"Iteration {iter_num}: x = {x}, grad = {grad}")
 
-x, result = polysolve.minimize(
+x, log = polysolve.minimize(
     Quadratic(),
     np.zeros(3),
     {
@@ -29,4 +29,4 @@ def post_step(self, iter_num, solver_info, x, grad):
 )
 
 print("Optimal point:", x)
-print("Optimal value:", result)
+print("Optimal value:", log)

src/polysolve.cpp

@@ -8,6 +8,7 @@
 #include <pybind11_json/pybind11_json.hpp>
 
 #include <polysolve/Types.hpp>
+#include <polysolve/linear/Solver.hpp>
 #include <polysolve/nonlinear/Problem.hpp>
 #include <polysolve/nonlinear/Solver.hpp>
 
@@ -27,6 +28,7 @@ namespace
     using TVector = Problem::TVector;
     using TMatrix = Problem::TMatrix;
     using THessian = Problem::THessian;
+    using StiffnessMatrix = polysolve::StiffnessMatrix;
 
     std::shared_ptr<spdlog::logger> python_logger(const spdlog::level::level_enum log_level)
     {
@@ -101,6 +103,55 @@ namespace
         }
     }
 
+    StiffnessMatrix cast_sparse_matrix(const py::object &obj)
+    {
+        StiffnessMatrix A;
+        try
+        {
+            A = obj.cast<StiffnessMatrix>();
+        }
+        catch (const py::cast_error &)
+        {
+            A = obj.cast<TMatrix>().sparseView();
+        }
+
+        if (A.rows() != A.cols())
+        {
+            throw std::runtime_error(
+                "linear system matrix must be square, got shape (" + std::to_string(A.rows()) + ", "
+                + std::to_string(A.cols()) + ")");
+        }
+
+        return A;
+    }
+
+    TMatrix cast_dense_matrix(const py::object &obj)
+    {
+        TMatrix A;
+        try
+        {
+            A = obj.cast<TMatrix>();
+        }
+        catch (const py::cast_error &)
+        {
+            A = TMatrix(obj.cast<StiffnessMatrix>());
+        }
+
+        if (A.rows() != A.cols())
+        {
+            throw std::runtime_error(
+                "linear system matrix must be square, got shape (" + std::to_string(A.rows()) + ", "
+                + std::to_string(A.cols()) + ")");
+        }
+
+        return A;
+    }
+
+    int default_precond_num(const Eigen::Index rows, const int precond_num)
+    {
+        return precond_num < 0 ? static_cast<int>(rows) : precond_num;
+    }
+
     class PyProblem : public Problem
     {
     public:
@@ -205,15 +256,131 @@ namespace
         }
     };
 
+    class LinearSolver
+    {
+    public:
+        LinearSolver(
+            const py::dict &solver_params,
+            const spdlog::level::level_enum log_level,
+            const bool strict_validation)
+            : logger_(python_logger(log_level))
+        {
+            py::scoped_ostream_redirect stdout_redirect(
+                std::cout, py::module_::import("sys").attr("stdout"));
+            const polysolve::json params = solver_params.cast<polysolve::json>();
+            solver_ = polysolve::linear::Solver::create(params, *logger_, strict_validation);
+        }
+
+        void set_parameters(const py::dict &params)
+        {
+            const polysolve::json jparams = params.cast<polysolve::json>();
+            solver_->set_parameters(jparams);
+        }
+
+        void analyze_pattern(const py::object &A_obj, const int precond_num)
+        {
+            py::scoped_ostream_redirect stdout_redirect(
+                std::cout, py::module_::import("sys").attr("stdout"));
+
+            if (solver_->is_dense())
+            {
+                TMatrix A = cast_dense_matrix(A_obj);
+                const int actual_precond_num = default_precond_num(A.rows(), precond_num);
+                py::gil_scoped_release release;
+                solver_->analyze_pattern_dense(A, actual_precond_num);
+            }
+            else
+            {
+                StiffnessMatrix A = cast_sparse_matrix(A_obj);
+                const int actual_precond_num = default_precond_num(A.rows(), precond_num);
+                py::gil_scoped_release release;
+                solver_->analyze_pattern(A, actual_precond_num);
+            }
+        }
+
+        void factorize(const py::object &A_obj)
+        {
+            py::scoped_ostream_redirect stdout_redirect(
+                std::cout, py::module_::import("sys").attr("stdout"));
+
+            if (solver_->is_dense())
+            {
+                TMatrix A = cast_dense_matrix(A_obj);
+                py::gil_scoped_release release;
+                solver_->factorize_dense(A);
+                factorized_size_ = A.rows();
+            }
+            else
+            {
+                StiffnessMatrix A = cast_sparse_matrix(A_obj);
+                py::gil_scoped_release release;
+                solver_->factorize(A);
+                factorized_size_ = A.rows();
+            }
+        }
+
+        TVector solve(const TVector &b, const py::object &x0_obj)
+        {
+            if (factorized_size_ < 0)
+                throw std::runtime_error("LinearSolver.factorize(A) must be called before solve(b)");
+            if (b.size() != factorized_size_)
+            {
+                throw std::runtime_error(
+                    "right-hand side has size " + std::to_string(b.size()) + ", expected "
+                    + std::to_string(factorized_size_));
+            }
+
+            TVector x;
+            if (x0_obj.is_none())
+            {
+                x = TVector::Zero(b.size());
+            }
+            else
+            {
+                x = x0_obj.cast<TVector>();
+                if (x.size() != b.size())
+                {
+                    throw std::runtime_error(
+                        "initial guess has size " + std::to_string(x.size()) + ", expected "
+                        + std::to_string(b.size()));
+                }
+            }
+
+            py::scoped_ostream_redirect stdout_redirect(
+                std::cout, py::module_::import("sys").attr("stdout"));
+            {
+                py::gil_scoped_release release;
+                solver_->solve(b, x);
+            }
+            return x;
+        }
+
+        polysolve::json info() const
+        {
+            polysolve::json result;
+            solver_->get_info(result);
+            return result;
+        }
+
+        std::string name() const { return solver_->name(); }
+        bool is_dense() const { return solver_->is_dense(); }
+
+    private:
+        std::shared_ptr<spdlog::logger> logger_;
+        std::unique_ptr<polysolve::linear::Solver> solver_;
+        Eigen::Index factorized_size_ = -1;
+    };
+
 } // namespace
 
 PYBIND11_MODULE(polysolve, m)
 {
     using namespace polysolve;
     namespace nonlinear = polysolve::nonlinear;
+    namespace linear = polysolve::linear;
     namespace nl = nlohmann;
 
-    m.doc() = "Python bindings for PolySolve nonlinear optimization.";
+    m.doc() = "Python bindings for PolySolve nonlinear optimization and linear solvers.";
 
     py::enum_<spdlog::level::level_enum>(m, "LogLevel")
         .value("trace", spdlog::level::trace)
@@ -240,6 +407,109 @@ PYBIND11_MODULE(polysolve, m)
             return hessian;
         });
 
+    py::class_<LinearSolver>(m, "LinearSolver")
+        .def(
+            py::init<const py::dict &, spdlog::level::level_enum, bool>(),
+            "Create a reusable PolySolve linear solver.",
+            py::arg("solver_params") = py::dict(),
+            py::arg("log_level") = spdlog::level::info,
+            py::arg("strict_validation") = true)
+        .def(
+            "set_parameters",
+            &LinearSolver::set_parameters,
+            "Set linear solver parameters.",
+            py::arg("params"))
+        .def(
+            "analyze_pattern",
+            &LinearSolver::analyze_pattern,
+            "Analyze the system matrix sparsity pattern.",
+            py::arg("A"),
+            py::arg("precond_num") = -1)
+        .def(
+            "analyse_pattern",
+            &LinearSolver::analyze_pattern,
+            "Analyze the system matrix sparsity pattern.",
+            py::arg("A"),
+            py::arg("precond_num") = -1)
+        .def(
+            "factorize",
+            &LinearSolver::factorize,
+            "Factorize the system matrix.",
+            py::arg("A"))
+        .def(
+            "factorise",
+            &LinearSolver::factorize,
+            "Factorize the system matrix.",
+            py::arg("A"))
+        .def(
+            "solve",
+            &LinearSolver::solve,
+            "Solve the factorized linear system.",
+            py::arg("b"),
+            py::arg("x0") = py::none())
+        .def(
+            "info",
+            &LinearSolver::info,
+            "Return information from the last solve.")
+        .def_property_readonly("name", &LinearSolver::name)
+        .def_property_readonly("is_dense", &LinearSolver::is_dense)
+        .def_static(
+            "available_solvers",
+            []() { return linear::Solver::available_solvers(); },
+            "List available linear solvers.")
+        .def_static(
+            "default_solver",
+            []() { return linear::Solver::default_solver(); },
+            "Return the default linear solver.")
+        .def_static(
+            "available_preconditioners",
+            []() { return linear::Solver::available_preconds(); },
+            "List available linear solver preconditioners.")
+        .def_static(
+            "default_preconditioner",
+            []() { return linear::Solver::default_precond(); },
+            "Return the default linear solver preconditioner.");
+
+    m.def(
+        "solve",
+        [](const py::object &A,
+           const TVector &b,
+           const py::dict &solver_params,
+           const py::object &x0,
+           const int precond_num,
+           const spdlog::level::level_enum log_level,
+           const bool strict_validation) {
+            LinearSolver solver(solver_params, log_level, strict_validation);
+            solver.analyze_pattern(A, precond_num);
+            solver.factorize(A);
+            return solver.solve(b, x0);
+        },
+        "Solve a linear system Ax = b.",
+        py::arg("A"),
+        py::arg("b"),
+        py::arg("solver_params") = py::dict(),
+        py::arg("x0") = py::none(),
+        py::arg("precond_num") = -1,
+        py::arg("log_level") = spdlog::level::info,
+        py::arg("strict_validation") = true);
+
+    m.def(
+        "available_linear_solvers",
+        []() { return linear::Solver::available_solvers(); },
+        "List available linear solvers.");
+    m.def(
+        "default_linear_solver",
+        []() { return linear::Solver::default_solver(); },
+        "Return the default linear solver.");
+    m.def(
+        "available_linear_preconditioners",
+        []() { return linear::Solver::available_preconds(); },
+        "List available linear solver preconditioners.");
+    m.def(
+        "default_linear_preconditioner",
+        []() { return linear::Solver::default_precond(); },
+        "Return the default linear solver preconditioner.");
+
     m.def(
         "minimize",
         [](Problem &problem,