Enable threading via Python interface

CMakeLists.txt

@@ -50,6 +50,7 @@ if(MPI_FOUND)
     message("-- MPI include files found in " "${MPI_INCLUDE_PATH}")
     separate_arguments(MPIEXEC_PREFLAGS) # to support multi flags
 endif()
+
 # OPENMP
 find_package(OpenMP)
 
@@ -83,7 +84,10 @@ if(MPI_FOUND)
     target_include_directories(Htool PRIVATE ${MPI_INCLUDE_PATH} ${MPI4PY_INCLUDE_DIR})
     target_compile_definitions(Htool PRIVATE HAVE_MPI)
 endif()
-target_link_libraries(Htool PRIVATE ${MPI_LIBRARIES} ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} ${ARPACK_LIBRARIES} ${OpenMP_CXX_LIBRARIES})
+target_link_libraries(Htool PRIVATE ${MPI_LIBRARIES} ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} ${ARPACK_LIBRARIES})
+if(OpenMP_CXX_FOUND)
+    target_link_libraries(Htool PRIVATE OpenMP::OpenMP_CXX)
+endif()
 
 if("${BLA_VENDOR}" STREQUAL "Intel10_32"
    OR "${BLA_VENDOR}" STREQUAL "Intel10_64lp"
@@ -94,7 +98,7 @@ if("${BLA_VENDOR}" STREQUAL "Intel10_32"
     target_compile_definitions(Htool PRIVATE "-DHPDDM_MKL -DHTOOL_MKL")
 endif()
 
-target_compile_definitions(Htool PRIVATE "-DHTOOL_WITH_PYTHON_INTERFACE" "-DHTOOL_WITH_HPDDM")
+target_compile_definitions(Htool PRIVATE "-DHTOOL_WITH_HPDDM")
 
 if(CODE_COVERAGE)
     if(CMAKE_C_COMPILER_ID MATCHES "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "GNU")

example/advanced/define_custom_low_rank_generator.py

@@ -29,3 +29,31 @@ def build_low_rank_approximation(self, rows, cols, epsilon):
         self.set_U(u[:, 0:truncated_rank] * s[0:truncated_rank])
         self.set_V(vh[0:truncated_rank, :])
         return True
+
+
+class ComplexCustomSVD(Htool.VirtualComplexLowRankGenerator):
+    def __init__(
+        self, generator: Htool.ComplexVirtualGenerator, allow_copy: bool = True
+    ):
+        super().__init__(allow_copy)
+        self.generator = generator
+
+    def build_low_rank_approximation(self, rows, cols, epsilon):
+        submat = np.zeros((len(rows), len(cols)), order="F", dtype=np.complex128)
+        self.generator.build_submatrix(rows, cols, submat)
+        u, s, vh = np.linalg.svd(submat, full_matrices=False)
+
+        norm = np.linalg.norm(submat)
+        svd_norm = 0
+        truncated_rank = len(s) - 1
+        while truncated_rank > 0 and math.sqrt(svd_norm) / norm < epsilon:
+            svd_norm += s[truncated_rank] ** 2
+            truncated_rank -= 1
+        truncated_rank += 1
+
+        if truncated_rank * (len(rows) + len(cols)) > (len(rows) * len(cols)):
+            return False  # the low rank approximation is not worthwhile
+
+        self.set_U(u[:, 0:truncated_rank] * s[0:truncated_rank])
+        self.set_V(vh[0:truncated_rank, :])
+        return True

example/define_generators.py

@@ -13,14 +13,14 @@ def __init__(self, target_points, source_points):
 
     def get_coef(self, i, j):
         return 1.0 / (
-            1e-1 + np.linalg.norm(self.target_points[:, i] - self.source_points[:, j])
+            1e-5 + np.linalg.norm(self.target_points[:, i] - self.source_points[:, j])
         )
 
     def build_submatrix(self, J, K, mat):
         for j in range(0, len(J)):
             for k in range(0, len(K)):
                 mat[j, k] = 1.0 / (
-                    1e-1
+                    1e-5
                     + np.linalg.norm(
                         self.target_points[:, J[j]] - self.source_points[:, K[k]]
                     )
@@ -41,3 +41,46 @@ def mat_mat(self, X):
                 for k in range(0, self.nb_cols):
                     Y[i, j] += self.get_coef(i, k) * X[k, j]
         return Y
+
+
+class ComplexCustomGenerator(Htool.ComplexVirtualGenerator):
+    def __init__(self, target_points, source_points):
+        super().__init__()
+        self.target_points = target_points
+        self.source_points = source_points
+        self.nb_rows = target_points.shape[1]
+        self.nb_cols = source_points.shape[1]
+
+    def get_coef(self, i, j):
+        diff = self.target_points[:, i] - self.source_points[:, j]
+        sign = 1 if diff[0] > 0 else -1
+        if i == j:
+            sign = 0
+        r = np.linalg.norm(diff)
+        return (1 + sign * 1j) / (1e-5 + r)
+
+    def build_submatrix(self, J, K, mat):
+        for j in range(0, len(J)):
+            for k in range(0, len(K)):
+                diff = self.target_points[:, J[j]] - self.source_points[:, K[k]]
+                sign = 1 if diff[0] > 0 else -1
+                if J[j] == K[k]:
+                    sign = 0
+                r = np.linalg.norm(diff)
+                mat[j, k] = (1 + sign * 1j) / (1e-5 + r)
+
+    def mat_vec(self, x):
+        y = np.zeros(self.nb_rows, dtype=np.complex128)
+        for i in range(0, self.nb_rows):
+            for j in range(0, self.nb_cols):
+                y[i] += self.get_coef(i, j) * x[j]
+        return y
+
+    def mat_mat(self, X):
+        Y = np.zeros((self.nb_rows, X.shape[1]), dtype=np.complex128)
+
+        for i in range(0, self.nb_rows):
+            for j in range(0, X.shape[1]):
+                for k in range(0, self.nb_cols):
+                    Y[i, j] += self.get_coef(i, k) * X[k, j]
+        return Y

pyproject.toml

@@ -2,7 +2,7 @@
 name = "Htool-DDM"
 version = "1.0.1rc2"
 authors = [{ name = "Pierre Marchand", email = "test@test.com" }]
-description = "CLI utility to automate asciinema"
+description = "Pybind11 interface for Htool-DDM"
 readme = "README.md"
 requires-python = ">=3.7"
 license = "MIT"
@@ -24,7 +24,8 @@ dev = ["ruff", "matplotlib>=3.0.0", "pytest", "scipy", "pybind11-stubgen"]
 
 
 [project.urls]
-"Homepage" = "https://github.com/htool-ddm/htool_python"
+"Homepage" = "https://htool-ddm.pages.math.cnrs.fr"
+"Source code" = "https://github.com/htool-ddm/htool_python"
 "Bug Tracker" = "https://github.com/htool-ddm/htool_python/issues"
 
 [tool.ruff]

setup.py

@@ -191,7 +191,7 @@ def build_extension(self, ext):
 
 
 setup(
-    ext_modules=[CMakeExtension("Htool-DDM")],
+    ext_modules=[CMakeExtension("Htool")],
     cmdclass=dict(build_ext=CMakeBuild),
     zip_safe=False,
     package_data={"": ["*.pyi"]},

src/htool/hmatrix/execution_policies.hpp

@@ -0,0 +1,52 @@
+#ifndef HTOOL_PYBIND11_EXECUTION_POLICIES_CPP
+#define HTOOL_PYBIND11_EXECUTION_POLICIES_CPP
+
+#include <htool/hmatrix/execution_policies.hpp>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+using namespace htool;
+
+inline void declare_standard_policies(py::module &m) {
+    py::class_<exec_compat::sequenced_policy>(
+        m,
+        "SequentialPolicy")
+        .def(py::init<>());
+
+    py::class_<exec_compat::parallel_policy>(
+        m,
+        "ParallelPolicy")
+        .def(py::init<>());
+}
+
+template <typename CoefficientPrecision, typename CoordinatePrecision>
+void declare_omp_task_policy(py::module &m, std::string prefix = "") {
+    py::class_<htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision>>(
+        m,
+        (prefix + "OmpTaskPolicy").c_str())
+        .def(py::init<>())
+        .def_property(
+            "max_number_of_nodes", [](htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &self) { return self.hmatrix_task_dependencies.max_number_of_nodes; }, [](htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &self, int value) { self.hmatrix_task_dependencies.max_number_of_nodes = value; })
+        .def(
+            "set_L0",
+            [](
+                htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &self,
+
+                HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix) {
+                self.hmatrix_task_dependencies.set_L0(hmatrix);
+            },
+            py::keep_alive<1, 2>())
+        .def_property_readonly(
+            "L0",
+            [](htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &self) {
+                py::list result;
+                for (HMatrix<CoefficientPrecision, CoordinatePrecision> *hmatrix : self.hmatrix_task_dependencies.L0) {
+                    result.append(py::cast(hmatrix, py::return_value_policy::reference));
+                }
+
+                return result;
+            },
+            py::return_value_policy::reference_internal);
+}
+
+#endif

src/htool/hmatrix/hmatrix.hpp

@@ -56,11 +56,33 @@ void declare_HMatrix(py::module &m, const std::string &className) {
     py_class.def("get_source_cluster", &HMatrix<CoefficientPrecision, CoordinatePrecision>::get_source_cluster, py::return_value_policy::reference_internal);
 
     py_class.def("lu_factorization", [](HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix) {
-        htool::lu_factorization(hmatrix);
+        htool::lu_factorization(exec_compat::par, hmatrix);
     });
+    py_class.def("lu_factorization", [](exec_compat::sequenced_policy &policy, HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix) {
+        htool::lu_factorization(exec_compat::seq, hmatrix);
+    });
+    py_class.def("lu_factorization", [](exec_compat::parallel_policy &policy, HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix) {
+        htool::lu_factorization(exec_compat::par, hmatrix);
+    });
+#ifdef _OPENMP
+    py_class.def("lu_factorization", [](htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &policy, HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix) {
+        htool::lu_factorization(policy, hmatrix);
+    });
+#endif
     py_class.def("cholesky_factorization", [](HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix, char UPLO) {
-        htool::cholesky_factorization(UPLO, hmatrix);
+        htool::cholesky_factorization(exec_compat::par, UPLO, hmatrix);
     });
+    py_class.def("cholesky_factorization", [](exec_compat::sequenced_policy &policy, HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix, char UPLO) {
+        htool::cholesky_factorization(exec_compat::seq, UPLO, hmatrix);
+    });
+    py_class.def("cholesky_factorization", [](exec_compat::parallel_policy &policy, HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix, char UPLO) {
+        htool::cholesky_factorization(exec_compat::par, UPLO, hmatrix);
+    });
+#ifdef _OPENMP
+    py_class.def("cholesky_factorization", [](htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &policy, HMatrix<CoefficientPrecision, CoordinatePrecision> &hmatrix, char UPLO) {
+        htool::cholesky_factorization(policy, UPLO, hmatrix);
+    });
+#endif
     py_class.def("lu_solve", [](const Class &self, char trans, const py::array_t<CoefficientPrecision, py::array::f_style> &input) {
         std::vector<ssize_t> shape;
         if (input.ndim() == 1) {
@@ -110,7 +132,7 @@ void declare_HMatrix(py::module &m, const std::string &className) {
             std::fill_n(result.mutable_data(), self.get_target_cluster().get_size(), CoefficientPrecision(0));
 
             char trans = 'N';
-            htool::add_hmatrix_vector_product(trans, CoefficientPrecision(1), self, input.data(), CoefficientPrecision(0), result.mutable_data());
+            htool::add_hmatrix_vector_product(exec_compat::par, trans, CoefficientPrecision(1), self, input.data(), CoefficientPrecision(0), result.mutable_data());
 
             return result;
         },
@@ -131,7 +153,7 @@ void declare_HMatrix(py::module &m, const std::string &className) {
             htool::MatrixView<CoefficientPrecision> output_view(input.shape()[0], input.shape()[1], result.mutable_data());
             char transa = 'N';
             char transb = 'N';
-            htool::add_hmatrix_matrix_product(transa, transb, CoefficientPrecision(1), self, input_view, CoefficientPrecision(0), output_view);
+            htool::add_hmatrix_matrix_product(exec_compat::par, transa, transb, CoefficientPrecision(1), self, input_view, CoefficientPrecision(0), output_view);
 
             return result;
         },

src/htool/hmatrix/hmatrix_tree_builder.hpp

@@ -33,7 +33,15 @@ void declare_hmatrix_builder(py::module &m, const std::string &className) {
     // LCOV_EXCL_STOP
 
     // Build
-    py_class.def("build", [](const Class &self, const VirtualGenerator<CoefficientPrecision> &generator, const Cluster<CoordinatePrecision> &target_cluster, const Cluster<CoordinatePrecision> &source_cluster, int target_partition_number, int partition_number_for_symmetry) { return self.build(generator, target_cluster, source_cluster, target_partition_number, partition_number_for_symmetry); }, py::arg("generator"), py::arg("target_cluster"), py::arg("source_cluster"), py::arg("target_partition_number") = -1, py::arg("partition_number_for_symmetry") = -1);
+    py_class.def("build", [](const Class &self, const exec_compat::sequenced_policy &policy, const VirtualGenerator<CoefficientPrecision> &generator, const Cluster<CoordinatePrecision> &target_cluster, const Cluster<CoordinatePrecision> &source_cluster, int target_partition_number, int partition_number_for_symmetry) { return self.build(exec_compat::seq, generator, target_cluster, source_cluster, target_partition_number, partition_number_for_symmetry); }, py::arg("policy"), py::arg("generator"), py::arg("target_cluster"), py::arg("source_cluster"), py::arg("target_partition_number") = -1, py::arg("partition_number_for_symmetry") = -1, py::call_guard<py::gil_scoped_release>());
+
+    py_class.def("build", [](const Class &self, const exec_compat::parallel_policy &policy, const VirtualGenerator<CoefficientPrecision> &generator, const Cluster<CoordinatePrecision> &target_cluster, const Cluster<CoordinatePrecision> &source_cluster, int target_partition_number, int partition_number_for_symmetry) { return self.build(exec_compat::par, generator, target_cluster, source_cluster, target_partition_number, partition_number_for_symmetry); }, py::arg("policy"), py::arg("generator"), py::arg("target_cluster"), py::arg("source_cluster"), py::arg("target_partition_number") = -1, py::arg("partition_number_for_symmetry") = -1, py::call_guard<py::gil_scoped_release>());
+
+#ifdef _OPENMP
+    py_class.def("build", [](const Class &self, htool::omp_task_policy<CoefficientPrecision, CoordinatePrecision> &policy, const VirtualGenerator<CoefficientPrecision> &generator, const Cluster<CoordinatePrecision> &target_cluster, const Cluster<CoordinatePrecision> &source_cluster, int target_partition_number, int partition_number_for_symmetry) { return self.build(policy, generator, target_cluster, source_cluster, target_partition_number, partition_number_for_symmetry); }, py::arg("policy"), py::arg("generator"), py::arg("target_cluster"), py::arg("source_cluster"), py::arg("target_partition_number") = -1, py::arg("partition_number_for_symmetry") = -1, py::call_guard<py::gil_scoped_release>());
+#endif
+
+    py_class.def("build", [](const Class &self, const VirtualGenerator<CoefficientPrecision> &generator, const Cluster<CoordinatePrecision> &target_cluster, const Cluster<CoordinatePrecision> &source_cluster, int target_partition_number, int partition_number_for_symmetry) { return self.build(exec_compat::par, generator, target_cluster, source_cluster, target_partition_number, partition_number_for_symmetry); }, py::arg("generator"), py::arg("target_cluster"), py::arg("source_cluster"), py::arg("target_partition_number") = -1, py::arg("partition_number_for_symmetry") = -1, py::call_guard<py::gil_scoped_release>());
 
     // Setters
     py_class.def("set_minimal_source_depth", &Class::set_minimal_source_depth);

src/htool/hmatrix/interfaces/virtual_generator.hpp

@@ -10,11 +10,11 @@ template <typename CoefficientPrecision>
 class VirtualGeneratorPython : public htool::VirtualGenerator<CoefficientPrecision> {
   public:
     using VirtualGenerator<CoefficientPrecision>::VirtualGenerator;
-
     VirtualGeneratorPython(const py::array_t<int> &target_permutation, const py::array_t<int> &source_permutation) : VirtualGenerator<CoefficientPrecision>() {}
 
     void copy_submatrix(int M, int N, const int *const rows, const int *const cols, CoefficientPrecision *ptr) const override {
         if (M * N > 0) {
+            py::gil_scoped_acquire acquire;
             py::array_t<CoefficientPrecision, py::array::f_style> mat(std::array<long int, 2>{M, N}, ptr, py::capsule(ptr));
 
             py::array_t<int> py_rows(std::array<long int, 1>{M}, rows, py::capsule(rows));
@@ -47,14 +47,60 @@ class PyVirtualGenerator : public VirtualGeneratorPython<CoefficientPrecision> {
 };
 
 template <typename CoefficientPrecision>
-void declare_virtual_generator(py::module &m, const std::string &className, const std::string &base_class_name) {
+class GeneratorFromCapsule : public htool::VirtualGenerator<CoefficientPrecision> {
+  private:
+    using backend_callback_t = void (*)(int, int, const int *, const int *, double *, void *);
+
+    struct Backend {
+        backend_callback_t callback;
+        void *ctx;
+    };
+
+    backend_callback_t callback = nullptr;
+    void *ctx                   = nullptr;
+
+    py::capsule m_capsule;
+
+  public:
+    using VirtualGenerator<CoefficientPrecision>::VirtualGenerator;
+    GeneratorFromCapsule(py::object obj) : VirtualGenerator<CoefficientPrecision>() {
+        if (!py::isinstance<py::capsule>(obj)) {
+            throw std::runtime_error("Backend needs to be a Python capsule");
+        }
+
+        m_capsule = obj.cast<py::capsule>();
+
+        auto *user_backend = static_cast<Backend *>(
+            PyCapsule_GetPointer(m_capsule.ptr(), "backend"));
+
+        if (!user_backend || !user_backend->callback) {
+            throw std::runtime_error("Invalid backend capsule");
+        }
+
+        callback = user_backend->callback;
+        ctx      = user_backend->ctx;
+    }
+
+    void copy_submatrix(int M, int N, const int *const rows, const int *const cols, CoefficientPrecision *ptr) const override {
+        double *user_ptr = reinterpret_cast<double *>(ptr);
+        if (M * N > 0) {
+            callback(M, N, rows, cols, user_ptr, ctx);
+        }
+    }
+};
+
+template <typename CoefficientPrecision>
+void declare_virtual_generator(py::module &m, const std::string &prefix) {
     using BaseClass = VirtualGenerator<CoefficientPrecision>;
-    py::class_<BaseClass>(m, base_class_name.c_str());
+    py::class_<BaseClass>(m, (prefix + "IGenerator").c_str());
 
     using Class = VirtualGeneratorPython<CoefficientPrecision>;
-    py::class_<Class, BaseClass, PyVirtualGenerator<CoefficientPrecision>> py_class(m, className.c_str());
+    py::class_<Class, BaseClass, PyVirtualGenerator<CoefficientPrecision>> py_class(m, (prefix + "VirtualGenerator").c_str());
     py_class.def(py::init<>());
     py_class.def("build_submatrix", &Class::build_submatrix);
+
+    py::class_<GeneratorFromCapsule<CoefficientPrecision>, BaseClass> py_class_2(m, (prefix + "GeneratorFromCapsule").c_str());
+    py_class_2.def(py::init<py::object>());
 }
 
 #endif

src/htool/hmatrix/interfaces/virtual_low_rank_generator.hpp

@@ -23,11 +23,11 @@ class VirtualLowRankGeneratorPython : public VirtualLowRankGenerator<Coefficient
     VirtualLowRankGeneratorPython(bool allow_copy = true) : VirtualLowRankGenerator<CoefficientPrecision>(), m_allow_copy(allow_copy) {}
 
     bool copy_low_rank_approximation(int M, int N, const int *const rows, const int *const cols, LowRankMatrix<CoefficientPrecision> &lrmat) const override {
+        py::gil_scoped_acquire acquire;
         auto &U = lrmat.get_U();
         auto &V = lrmat.get_V();
         py::array_t<int> py_rows(std::array<long int, 1>{M}, rows, py::capsule(rows));
         py::array_t<int> py_cols(std::array<long int, 1>{N}, cols, py::capsule(cols));
-
         bool success = build_low_rank_approximation(py_rows, py_cols, lrmat.get_epsilon());
         if (success) {
             if (m_allow_copy) {