diff --git a/CMakeLists.txt b/CMakeLists.txt index acd125e..4bfe877 100644 --- a/CMakeLists.txt +++ b/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") diff --git a/example/advanced/define_custom_low_rank_generator.py b/example/advanced/define_custom_low_rank_generator.py index f9ed9e7..58b1ce6 100644 --- a/example/advanced/define_custom_low_rank_generator.py +++ b/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 diff --git a/example/define_generators.py b/example/define_generators.py index 9349cce..36d01da 100644 --- a/example/define_generators.py +++ b/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 diff --git a/lib/htool b/lib/htool index f0b1541..e292614 160000 --- a/lib/htool +++ b/lib/htool @@ -1 +1 @@ -Subproject commit f0b154130a0f542b19ad6d3b4e6986539fb702b7 +Subproject commit e292614ce682e27b3c65b18ed1e27cbd4ac1923d diff --git a/lib/pybind11 b/lib/pybind11 index cc86e8b..678b673 160000 --- a/lib/pybind11 +++ b/lib/pybind11 @@ -1 +1 @@ -Subproject commit cc86e8b2a6c68580f23c433479e2815aefa4d880 +Subproject commit 678b6735e44fe9d7feac7ff31bc9c86d42c04613 diff --git a/pyproject.toml b/pyproject.toml index db97704..e97aede 100644 --- a/pyproject.toml +++ b/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] diff --git a/setup.py b/setup.py index b473c1e..2739793 100644 --- a/setup.py +++ b/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"]}, diff --git a/src/htool/hmatrix/execution_policies.hpp b/src/htool/hmatrix/execution_policies.hpp new file mode 100644 index 0000000..6ccfb0b --- /dev/null +++ b/src/htool/hmatrix/execution_policies.hpp @@ -0,0 +1,52 @@ +#ifndef HTOOL_PYBIND11_EXECUTION_POLICIES_CPP +#define HTOOL_PYBIND11_EXECUTION_POLICIES_CPP + +#include +#include + +namespace py = pybind11; +using namespace htool; + +inline void declare_standard_policies(py::module &m) { + py::class_( + m, + "SequentialPolicy") + .def(py::init<>()); + + py::class_( + m, + "ParallelPolicy") + .def(py::init<>()); +} + +template +void declare_omp_task_policy(py::module &m, std::string prefix = "") { + py::class_>( + m, + (prefix + "OmpTaskPolicy").c_str()) + .def(py::init<>()) + .def_property( + "max_number_of_nodes", [](htool::omp_task_policy &self) { return self.hmatrix_task_dependencies.max_number_of_nodes; }, [](htool::omp_task_policy &self, int value) { self.hmatrix_task_dependencies.max_number_of_nodes = value; }) + .def( + "set_L0", + []( + htool::omp_task_policy &self, + + HMatrix &hmatrix) { + self.hmatrix_task_dependencies.set_L0(hmatrix); + }, + py::keep_alive<1, 2>()) + .def_property_readonly( + "L0", + [](htool::omp_task_policy &self) { + py::list result; + for (HMatrix *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 diff --git a/src/htool/hmatrix/hmatrix.hpp b/src/htool/hmatrix/hmatrix.hpp index 27d7752..291203b 100644 --- a/src/htool/hmatrix/hmatrix.hpp +++ b/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::get_source_cluster, py::return_value_policy::reference_internal); py_class.def("lu_factorization", [](HMatrix &hmatrix) { - htool::lu_factorization(hmatrix); + htool::lu_factorization(exec_compat::par, hmatrix); }); + py_class.def("lu_factorization", [](exec_compat::sequenced_policy &policy, HMatrix &hmatrix) { + htool::lu_factorization(exec_compat::seq, hmatrix); + }); + py_class.def("lu_factorization", [](exec_compat::parallel_policy &policy, HMatrix &hmatrix) { + htool::lu_factorization(exec_compat::par, hmatrix); + }); +#ifdef _OPENMP + py_class.def("lu_factorization", [](htool::omp_task_policy &policy, HMatrix &hmatrix) { + htool::lu_factorization(policy, hmatrix); + }); +#endif py_class.def("cholesky_factorization", [](HMatrix &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 &hmatrix, char UPLO) { + htool::cholesky_factorization(exec_compat::seq, UPLO, hmatrix); + }); + py_class.def("cholesky_factorization", [](exec_compat::parallel_policy &policy, HMatrix &hmatrix, char UPLO) { + htool::cholesky_factorization(exec_compat::par, UPLO, hmatrix); + }); +#ifdef _OPENMP + py_class.def("cholesky_factorization", [](htool::omp_task_policy &policy, HMatrix &hmatrix, char UPLO) { + htool::cholesky_factorization(policy, UPLO, hmatrix); + }); +#endif py_class.def("lu_solve", [](const Class &self, char trans, const py::array_t &input) { std::vector 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 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; }, diff --git a/src/htool/hmatrix/hmatrix_tree_builder.hpp b/src/htool/hmatrix/hmatrix_tree_builder.hpp index 217fe28..a787f31 100644 --- a/src/htool/hmatrix/hmatrix_tree_builder.hpp +++ b/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 &generator, const Cluster &target_cluster, const Cluster &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 &generator, const Cluster &target_cluster, const Cluster &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_class.def("build", [](const Class &self, const exec_compat::parallel_policy &policy, const VirtualGenerator &generator, const Cluster &target_cluster, const Cluster &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()); + +#ifdef _OPENMP + py_class.def("build", [](const Class &self, htool::omp_task_policy &policy, const VirtualGenerator &generator, const Cluster &target_cluster, const Cluster &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()); +#endif + + py_class.def("build", [](const Class &self, const VirtualGenerator &generator, const Cluster &target_cluster, const Cluster &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()); // Setters py_class.def("set_minimal_source_depth", &Class::set_minimal_source_depth); diff --git a/src/htool/hmatrix/interfaces/virtual_generator.hpp b/src/htool/hmatrix/interfaces/virtual_generator.hpp index 532c160..011606b 100644 --- a/src/htool/hmatrix/interfaces/virtual_generator.hpp +++ b/src/htool/hmatrix/interfaces/virtual_generator.hpp @@ -10,11 +10,11 @@ template class VirtualGeneratorPython : public htool::VirtualGenerator { public: using VirtualGenerator::VirtualGenerator; - VirtualGeneratorPython(const py::array_t &target_permutation, const py::array_t &source_permutation) : VirtualGenerator() {} 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 mat(std::array{M, N}, ptr, py::capsule(ptr)); py::array_t py_rows(std::array{M}, rows, py::capsule(rows)); @@ -47,14 +47,60 @@ class PyVirtualGenerator : public VirtualGeneratorPython { }; template -void declare_virtual_generator(py::module &m, const std::string &className, const std::string &base_class_name) { +class GeneratorFromCapsule : public htool::VirtualGenerator { + 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::VirtualGenerator; + GeneratorFromCapsule(py::object obj) : VirtualGenerator() { + if (!py::isinstance(obj)) { + throw std::runtime_error("Backend needs to be a Python capsule"); + } + + m_capsule = obj.cast(); + + auto *user_backend = static_cast( + 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(ptr); + if (M * N > 0) { + callback(M, N, rows, cols, user_ptr, ctx); + } + } +}; + +template +void declare_virtual_generator(py::module &m, const std::string &prefix) { using BaseClass = VirtualGenerator; - py::class_(m, base_class_name.c_str()); + py::class_(m, (prefix + "IGenerator").c_str()); using Class = VirtualGeneratorPython; - py::class_> py_class(m, className.c_str()); + py::class_> py_class(m, (prefix + "VirtualGenerator").c_str()); py_class.def(py::init<>()); py_class.def("build_submatrix", &Class::build_submatrix); + + py::class_, BaseClass> py_class_2(m, (prefix + "GeneratorFromCapsule").c_str()); + py_class_2.def(py::init()); } #endif diff --git a/src/htool/hmatrix/interfaces/virtual_low_rank_generator.hpp b/src/htool/hmatrix/interfaces/virtual_low_rank_generator.hpp index d0f1a21..c47ae65 100644 --- a/src/htool/hmatrix/interfaces/virtual_low_rank_generator.hpp +++ b/src/htool/hmatrix/interfaces/virtual_low_rank_generator.hpp @@ -23,11 +23,11 @@ class VirtualLowRankGeneratorPython : public VirtualLowRankGenerator(), m_allow_copy(allow_copy) {} bool copy_low_rank_approximation(int M, int N, const int *const rows, const int *const cols, LowRankMatrix &lrmat) const override { + py::gil_scoped_acquire acquire; auto &U = lrmat.get_U(); auto &V = lrmat.get_V(); py::array_t py_rows(std::array{M}, rows, py::capsule(rows)); py::array_t py_cols(std::array{N}, cols, py::capsule(cols)); - bool success = build_low_rank_approximation(py_rows, py_cols, lrmat.get_epsilon()); if (success) { if (m_allow_copy) { diff --git a/src/htool/main.cpp b/src/htool/main.cpp index 548445d..64e3151 100644 --- a/src/htool/main.cpp +++ b/src/htool/main.cpp @@ -8,6 +8,7 @@ #include "clustering/interface/virtual_partitioning.hpp" #include "clustering/utility.hpp" +#include "hmatrix/execution_policies.hpp" #include "hmatrix/hmatrix.hpp" #include "hmatrix/hmatrix_tree_builder.hpp" #include "hmatrix/interfaces/virtual_dense_blocks_generator.hpp" @@ -58,11 +59,15 @@ PYBIND11_MODULE(Htool, m) { declare_cluster_builder(m, "ClusterTreeBuilder"); declare_cluster_utility(m); - declare_virtual_generator(m, "VirtualGenerator", "IGenerator"); + declare_virtual_generator(m, ""); declare_LowRankMatrix(m, "LowRankMatrix"); declare_HMatrix(m, "HMatrix"); declare_custom_VirtualLowRankGenerator(m, "VirtualLowRankGenerator"); declare_custom_VirtualDenseBlocksGenerator(m, "VirtualDenseBlocksGenerator"); + declare_standard_policies(m); +#ifdef _OPENMP + declare_omp_task_policy(m); +#endif declare_hmatrix_builder(m, "HMatrixTreeBuilder"); #ifdef HAVE_MPI @@ -89,9 +94,12 @@ PYBIND11_MODULE(Htool, m) { declare_virtual_partitioning>(m, "Complex"); declare_LowRankMatrix>(m, "ComplexLowRankMatrix"); declare_HMatrix, double>(m, "ComplexHMatrix"); - declare_virtual_generator>(m, "ComplexVirtualGenerator", "IComplexGenerator"); + declare_virtual_generator>(m, "Complex"); declare_custom_VirtualLowRankGenerator>(m, "VirtualComplexLowRankGenerator"); declare_custom_VirtualDenseBlocksGenerator>(m, "ComplexVirtualDenseBlocksGenerator"); +#ifdef _OPENMP + declare_omp_task_policy, double>(m, "Complex"); +#endif declare_hmatrix_builder, double>(m, "ComplexHMatrixTreeBuilder"); #ifdef HAVE_MPI diff --git a/src/htool/matplotlib/hmatrix.hpp b/src/htool/matplotlib/hmatrix.hpp index 75a42e9..1bdcce3 100644 --- a/src/htool/matplotlib/hmatrix.hpp +++ b/src/htool/matplotlib/hmatrix.hpp @@ -7,7 +7,7 @@ template void declare_matplotlib_hmatrix(py::module &m) { - m.def("plot", [](py::object axes, const HMatrix &hmatrix) { + m.def("plot", [](py::object axes, const HMatrix &hmatrix, py::object L0) { std::vector buf; int nb_leaves = 0; htool::preorder_tree_traversal( @@ -24,31 +24,25 @@ void declare_matplotlib_hmatrix(py::module &m) { }); // Import - py::object plt = py::module::import("matplotlib.pyplot"); - py::object patches = py::module::import("matplotlib.patches"); - py::object colors = py::module::import("matplotlib.colors"); - py::object numpy = py::module::import("numpy"); + py::object plt = py::module::import("matplotlib.pyplot"); + py::object patches = py::module::import("matplotlib.patches"); + py::object collections = py::module::import("matplotlib.collections"); + py::object colors = py::module::import("matplotlib.colors"); + py::object numpy = py::module::import("numpy"); + + // Colormap + py::object cmap = plt.attr("get_cmap")("YlGn"); + py::object new_cmap = colors.attr("LinearSegmentedColormap").attr("from_list")("trunc(YlGn,0.4,1)", cmap(numpy.attr("linspace")(0.4, 1, 100))); + py::object norm = colors.attr("Normalize")("vmin"_a = 0, "vmax"_a = 10); // First Data int nr = hmatrix.get_target_cluster().get_size(); int nc = hmatrix.get_source_cluster().get_size(); - py::array_t matrix({nr, nc}); - py::array_t mask_matrix({nr, nc}); - matrix.attr("fill")(0); - mask_matrix.attr("fill")(false); - - // Figure - // py::tuple sublots_output = plt.attr("subplots")(1, 1); - // py::object fig = sublots_output[0]; - // py::object axes = sublots_output[1]; - // axes.attr() - - // Issue: there a shift of one pixel along the y-axis... - // int shift = axes.transData.transform([(0,0), (1,1)]) - // shift = shift[1,1] - shift[0,1] # 1 unit in display coords - int shift = 0; - - int max_rank = 0; + + // Storage for rectangles and colors + py::list rects; + py::list facecolors; + for (int p = 0; p < nb_leaves; p++) { int i_row = buf[5 * p]; int nb_row = buf[5 * p + 1]; @@ -56,37 +50,91 @@ void declare_matplotlib_hmatrix(py::module &m) { int nb_col = buf[5 * p + 3]; int rank = buf[5 * p + 4]; - if (rank > max_rank) { - max_rank = rank; - } - for (int i = 0; i < nb_row; i++) { - for (int j = 0; j < nb_col; j++) { - matrix.mutable_at(i_row + i, i_col + j) = rank; - if (rank == -1) { - mask_matrix.mutable_at(i_row + i, i_col + j) = true; - } - } + // Color selection + py::object facecolor; + if (rank == -1) { + facecolors.append(py::str("red")); // full blocks + } else { + facecolors.append(new_cmap(norm(rank))); } - py::object rect = patches.attr("Rectangle")(py::make_tuple(i_col - 0.5, i_row - 0.5 + shift), nb_col, nb_row, "linewidth"_a = 0.75, "edgecolor"_a = 'k', "facecolor"_a = "none"); - axes.attr("add_patch")(rect); + // Rectangle (no styling here!) + py::object rect = patches.attr("Rectangle")( + py::make_tuple(i_col, i_row), + nb_col, + nb_row); + rects.append(rect); - if (rank >= 0 && nb_col / double(nc) > 0.05 && nb_row / double(nc) > 0.05) { - axes.attr("annotate")(rank, py::make_tuple(i_col + nb_col / 2., i_row + nb_row / 2.), "color"_a = "white", "size"_a = 10, "va"_a = "center", "ha"_a = "center"); + // Optional: annotate only large blocks + if (rank >= 0 && nb_col > nc * 0.05 && nb_row > nr * 0.05) { + axes.attr("text")( + i_col + nb_col / 2.0, + i_row + nb_row / 2.0, + rank, + "color"_a = "white", + "ha"_a = "center", + "va"_a = "center", + "fontsize"_a = 8); } } - // Colormap - py::object cmap = plt.attr("get_cmap")("YlGn"); - py::object new_cmap = colors.attr("LinearSegmentedColormap").attr("from_list")("trunc(YlGn,0.4,1)", cmap(numpy.attr("linspace")(0.4, 1, 100))); + // Create PatchCollection + py::object collection = collections.attr("PatchCollection")( + rects, + "facecolor"_a = facecolors, + "edgecolor"_a = "black", // huge speedup + "linewidth"_a = 0.2); + + axes.attr("add_collection")(collection); + + // Axes formatting + axes.attr("set_xlim")(0, nc); + axes.attr("set_ylim")(nr, 0); // invert y-axis (matrix style) + axes.attr("set_aspect")("equal"); + + // Optional: remove ticks for speed + axes.attr("set_xticks")(py::list()); + axes.attr("set_yticks")(py::list()); + + if (!L0.is_none()) { + py::list L0_rects; + + for (py::handle item : L0) { + const HMatrix &block = + item.cast &>(); + + const int i_row = + block.get_target_cluster().get_offset() + - hmatrix.get_target_cluster().get_offset(); + + const int nb_row = + block.get_target_cluster().get_size(); + + const int i_col = + block.get_source_cluster().get_offset() + - hmatrix.get_source_cluster().get_offset(); + + const int nb_col = + block.get_source_cluster().get_size(); + + py::object rect = + patches.attr("Rectangle")( + py::make_tuple(i_col, i_row), + nb_col, + nb_row); + + L0_rects.append(rect); + } - // Plot - py::object masked_matrix = numpy.attr("ma").attr("array")(matrix, "mask"_a = mask_matrix); - new_cmap.attr("set_bad")("color"_a = "red"); + py::object L0_collection = + collections.attr("PatchCollection")( + L0_rects, + "facecolor"_a = "none", + "edgecolor"_a = "purple", + "linewidth"_a = 1.5); - axes.attr("imshow")(masked_matrix, "cmap"_a = new_cmap, "vmin"_a = 0, "vmax"_a = 10); - // plt.attr("draw")(); - }); + axes.attr("add_collection")(L0_collection); + } }, py::arg("axes"), py::arg("hmatrix"), py::arg("L0") = py::none()); } #endif diff --git a/tests/test_hmatrix.py b/tests/test_hmatrix.py index 7eed45a..a73e8d3 100644 --- a/tests/test_hmatrix.py +++ b/tests/test_hmatrix.py @@ -5,25 +5,40 @@ import pytest import Htool -from example.advanced.define_custom_low_rank_generator import CustomSVD +from example.advanced.define_custom_low_rank_generator import ( + ComplexCustomSVD, + CustomSVD, +) from example.create_geometry import create_random_geometries -from example.define_generators import CustomGenerator +from example.define_generators import ComplexCustomGenerator, CustomGenerator @pytest.mark.parametrize( - "loglevel,symmetry", + "loglevel,symmetry,is_complex,policy_name", [ - (logging.INFO, "N"), - (logging.DEBUG, "N"), - (logging.WARNING, "N"), - (logging.ERROR, "N"), - (logging.CRITICAL, "N"), - (logging.INFO, "S"), + (logging.INFO, "N", False, None), + (logging.DEBUG, "N", False, None), + (logging.WARNING, "N", False, None), + (logging.ERROR, "N", False, None), + (logging.CRITICAL, "N", False, None), + (logging.INFO, "S", False, None), + (logging.INFO, "N", False, "par"), + (logging.INFO, "N", False, "seq"), + (logging.INFO, "N", False, "omp_task"), + (logging.INFO, "S", True, None), ], ) -def test_hmatrix(loglevel, symmetry): +def test_hmatrix(loglevel, symmetry, is_complex, policy_name): logging.basicConfig(level=loglevel) + policy = None + if policy_name == "par": + policy = Htool.ParallelPolicy() + elif policy_name == "seq": + policy = Htool.SequentialPolicy() + elif policy_name == "omp_task": + policy = Htool.ComplexOmpTaskPolicy() if is_complex else Htool.OmpTaskPolicy() + # Random geometry nb_rows = 500 nb_cols = 500 @@ -52,18 +67,37 @@ def test_hmatrix(loglevel, symmetry): source_cluster = target_cluster # Build generator - if symmetry == "N": - generator = CustomGenerator(target_points, source_points) + if is_complex is False: + if symmetry == "N": + generator = CustomGenerator(target_points, source_points) + else: + generator = CustomGenerator(target_points, target_points) else: - generator = CustomGenerator(target_points, target_points) - + if symmetry == "N": + generator = ComplexCustomGenerator(target_points, source_points) + else: + generator = ComplexCustomGenerator(target_points, target_points) # Custom low rank generator - low_rank_generator = CustomSVD(generator, False) + if is_complex is False: + low_rank_generator = CustomSVD(generator, False) + else: + low_rank_generator = ComplexCustomSVD(generator, False) # Build HMatrix - hmatrix_builder = Htool.HMatrixTreeBuilder(epsilon, eta, "N", "N") + if is_complex is False: + hmatrix_builder = Htool.HMatrixTreeBuilder(epsilon, eta, "N", "N") + else: + hmatrix_builder = Htool.ComplexHMatrixTreeBuilder(epsilon, eta, "N", "N") hmatrix_builder.set_low_rank_generator(low_rank_generator) - hmatrix = hmatrix_builder.build(generator, target_cluster, source_cluster) + if symmetry == "S": + hmatrix_builder.set_block_tree_consistency(True) + + if policy: + hmatrix = hmatrix_builder.build( + policy, generator, target_cluster, source_cluster + ) + else: + hmatrix = hmatrix_builder.build(generator, target_cluster, source_cluster) assert hmatrix.shape == (nb_rows, nb_cols) # Copy @@ -74,55 +108,68 @@ def test_hmatrix(loglevel, symmetry): dense_in_user_numbering = hmatrix.to_dense_in_user_numbering() # HMatrix vector product + dtype = np.float64 if is_complex is False else np.complex128 np.random.seed(0) - x = np.random.rand(nb_cols) + if is_complex is False: + x = np.random.rand(nb_cols) + else: + x = np.random.rand(nb_cols) + 1j * np.random.rand(nb_cols) y = hmatrix * x y_exact = generator.mat_vec(x) y_dense = dense_in_user_numbering.dot(x) y_copy = copy_hmatrix * x assert np.linalg.norm(y - y_exact) / np.linalg.norm(y_exact) < epsilon assert np.linalg.norm(y - y_dense) / np.linalg.norm(y_dense) < 1e-10 - assert np.linalg.norm(y - y_copy) < 1e-10 + assert np.linalg.norm(y - y_copy) / np.linalg.norm(y) < 1e-10 # HMatrix matrix product np.random.seed(0) - x = np.random.rand(nb_cols, 2) + if is_complex is False: + x = np.random.rand(nb_cols, 2) + else: + x = np.random.rand(nb_cols, 2) + 1j * np.random.rand(nb_cols, 2) y = hmatrix @ x y_exact = generator.mat_mat(x) y_dense = dense_in_user_numbering @ x y_copy = copy_hmatrix @ x assert np.linalg.norm(y - y_exact) / np.linalg.norm(y_exact) < epsilon assert np.linalg.norm(y - y_dense) / np.linalg.norm(y_dense) < 1e-10 - assert np.linalg.norm(y - y_copy) < 1e-10 + assert np.linalg.norm(y - y_copy) / np.linalg.norm(y) < 1e-10 if symmetry != "N": # HLU factorization - copy_hmatrix.lu_factorization() + if policy: + copy_hmatrix.lu_factorization(policy) + else: + copy_hmatrix.lu_factorization() # HLU solve vec - x_ref = np.ones(nb_cols) + x_ref = np.ones(nb_cols, dtype=dtype) y = hmatrix * x_ref x = copy_hmatrix.lu_solve("N", y) assert np.linalg.norm(x - x_ref) / np.linalg.norm(x_ref) < epsilon # HLU solve mat - x_ref = np.ones((nb_cols, 2)) + x_ref = np.ones((nb_cols, 2), dtype=dtype) y = hmatrix @ x_ref x = copy_hmatrix.lu_solve("N", y) assert np.linalg.norm(x - x_ref) / np.linalg.norm(x_ref) < epsilon # Cholesky factorization copy_hmatrix = copy.deepcopy(hmatrix) - copy_hmatrix.cholesky_factorization("L") + if policy: + copy_hmatrix.cholesky_factorization(policy, "L") + else: + copy_hmatrix.cholesky_factorization("L") # Cholesky solve vec - x_ref = np.ones(nb_cols) + x_ref = np.ones(nb_cols, dtype=dtype) y = hmatrix * x_ref x = copy_hmatrix.cholesky_solve("L", y) assert np.linalg.norm(x - x_ref) / np.linalg.norm(x_ref) < epsilon # Cholesky solve mat - x_ref = np.ones((nb_cols, 2)) + x_ref = np.ones((nb_cols, 2), dtype=dtype) y = hmatrix @ x_ref x = copy_hmatrix.cholesky_solve("L", y) assert np.linalg.norm(x - x_ref) / np.linalg.norm(x_ref) < epsilon