diff --git a/example/define_custom_dense_blocks_generator.py b/example/define_custom_dense_blocks_generator.py index 1778d0c..d72cbac 100644 --- a/example/define_custom_dense_blocks_generator.py +++ b/example/define_custom_dense_blocks_generator.py @@ -3,10 +3,9 @@ class CustomDenseBlocksGenerator(Htool.VirtualDenseBlocksGenerator): def __init__( - self, - generator, + self, generator, target_cluster: Htool.Cluster, source_cluster: Htool.Cluster ): - super().__init__() + super().__init__(target_cluster, source_cluster) self.generator = generator def build_dense_blocks(self, rows_offsets, cols_offsets, blocks): diff --git a/example/use_custom_dense_block_generator.py b/example/use_custom_dense_block_generator.py index 05639ba..2988c7d 100644 --- a/example/use_custom_dense_block_generator.py +++ b/example/use_custom_dense_block_generator.py @@ -1,17 +1,14 @@ import Htool import mpi4py import numpy as np -from create_geometry import create_partitionned_geometries +from create_geometry import create_random_geometries from define_custom_dense_blocks_generator import CustomDenseBlocksGenerator from define_custom_generators import CustomGenerator # Random geometry -nb_rows = 500 -nb_cols = 500 +size = 500 dimension = 3 -[target_points, source_points, target_partition] = create_partitionned_geometries( - dimension, nb_rows, nb_cols, mpi4py.MPI.COMM_WORLD.size -) +[points, _] = create_random_geometries(dimension, size, size) # Htool parameters @@ -24,69 +21,51 @@ cluster_builder = Htool.ClusterBuilder() cluster_builder.set_minclustersize(minclustersize) target_cluster: Htool.Cluster = cluster_builder.create_cluster_tree( - target_points, number_of_children, mpi4py.MPI.COMM_WORLD.size, target_partition -) -source_cluster: Htool.Cluster = cluster_builder.create_cluster_tree( - source_points, number_of_children, mpi4py.MPI.COMM_WORLD.size + points, number_of_children, mpi4py.MPI.COMM_WORLD.size ) # Build generator -generator = CustomGenerator(target_points, source_points) +generator = CustomGenerator(points, points) # -custom_dense_blocks_generator = CustomDenseBlocksGenerator(generator) +custom_dense_blocks_generator = CustomDenseBlocksGenerator( + generator, target_cluster, target_cluster +) # Build HMatrix hmatrix_builder = Htool.HMatrixBuilder( target_cluster, - source_cluster, + target_cluster, epsilon, eta, "N", "N", -1, mpi4py.MPI.COMM_WORLD.rank, + mpi4py.MPI.COMM_WORLD.rank, ) hmatrix_builder.set_dense_blocks_generator(custom_dense_blocks_generator) -hmatrix: Htool.HMatrix = hmatrix_builder.build(generator) - -# Build local operator -local_operator = Htool.LocalHMatrix( - hmatrix, - target_cluster.get_cluster_on_partition(mpi4py.MPI.COMM_WORLD.rank), - source_cluster, - "N", - "N", - False, - False, -) # Build distributed operator -target_partition_from_cluster = Htool.PartitionFromCluster(target_cluster) -source_partition_from_cluster = Htool.PartitionFromCluster(source_cluster) -distributed_operator = Htool.DistributedOperator( - target_partition_from_cluster, - source_partition_from_cluster, - "N", - "N", - mpi4py.MPI.COMM_WORLD, +distributed_operator_from_hmatrix = Htool.DistributedOperatorFromHMatrix( + generator, target_cluster, target_cluster, hmatrix_builder, mpi4py.MPI.COMM_WORLD ) -distributed_operator.add_local_operator(local_operator) - +distributed_operator = distributed_operator_from_hmatrix.distributed_operator +hmatrix = distributed_operator_from_hmatrix.hmatrix # Test matrix vector product np.random.seed(0) -x = np.random.rand(nb_cols) +x = np.random.rand(size) y_1 = distributed_operator * x y_2 = generator.mat_vec(x) print(mpi4py.MPI.COMM_WORLD.rank, np.linalg.norm(y_1 - y_2) / np.linalg.norm(y_2)) # Test matrix matrix product -X = np.asfortranarray(np.random.rand(nb_cols, 2)) +X = np.asfortranarray(np.random.rand(size, 2)) Y_1 = distributed_operator @ X Y_2 = generator.mat_mat(X) print(mpi4py.MPI.COMM_WORLD.rank, np.linalg.norm(Y_1 - Y_2) / np.linalg.norm(Y_2)) diff --git a/src/htool/hmatrix/interfaces/virtual_dense_blocks_generator.hpp b/src/htool/hmatrix/interfaces/virtual_dense_blocks_generator.hpp index 51a5e6b..d544e89 100644 --- a/src/htool/hmatrix/interfaces/virtual_dense_blocks_generator.hpp +++ b/src/htool/hmatrix/interfaces/virtual_dense_blocks_generator.hpp @@ -7,27 +7,35 @@ namespace py = pybind11; using namespace pybind11::literals; using namespace htool; -template +template > class VirtualDenseBlocksGeneratorPython : public VirtualDenseBlocksGenerator { + const Cluster &m_target_cluster; + const Cluster &m_source_cluster; + public: using VirtualDenseBlocksGenerator::VirtualDenseBlocksGenerator; - void copy_dense_blocks(const std::vector &M, const std::vector &N, const std::vector &rows_offsets, const std::vector &cols_offsets, std::vector &ptr) const override { + VirtualDenseBlocksGeneratorPython(const Cluster &target_cluster, const Cluster &source_cluster) : VirtualDenseBlocksGenerator(), m_target_cluster(target_cluster), m_source_cluster(source_cluster) {} - int nb_blocks = M.size(); - py::array_t rows_offsets_np(rows_offsets.size(), rows_offsets.data(), py::capsule(rows_offsets.data())); - py::array_t cols_offsets_np(cols_offsets.size(), cols_offsets.data(), py::capsule(cols_offsets.data())); + void copy_dense_blocks(const std::vector &M, const std::vector &N, const std::vector &rows_offsets, const std::vector &cols_offsets, std::vector &ptr) const override { + int nb_blocks = M.size(); + auto &target_permutation = m_target_cluster.get_permutation(); + auto &source_permutation = m_source_cluster.get_permutation(); std::vector> vec_ptr; + std::vector> rows_ptr; + std::vector> cols_ptr; for (int i = 0; i < nb_blocks; i++) { + rows_ptr.emplace_back(std::array{M[i]}, target_permutation.data() + rows_offsets[i], py::capsule(target_permutation.data() + rows_offsets[i])); + cols_ptr.emplace_back(std::array{N[i]}, source_permutation.data() + cols_offsets[i], py::capsule(source_permutation.data() + cols_offsets[i])); vec_ptr.emplace_back(std::array{M[i], N[i]}, ptr[i], py::capsule(ptr[i])); } - build_dense_blocks(rows_offsets_np, cols_offsets_np, vec_ptr); + build_dense_blocks(rows_ptr, cols_ptr, vec_ptr); } // lcov does not see it because of trampoline I assume - virtual void build_dense_blocks(const py::array_t &rows_offsets_np, const py::array_t &cols_offsets_np, std::vector> &blocks) const = 0; // LCOV_EXCL_LINE + virtual void build_dense_blocks(const std::vector> &rows, const std::vector> &cols, std::vector> &blocks) const = 0; // LCOV_EXCL_LINE }; template @@ -37,7 +45,7 @@ class PyVirtualDenseBlocksGenerator : public VirtualDenseBlocksGeneratorPython(nr0,nc0){} /* Trampoline (need one for each virtual function) */ - virtual void build_dense_blocks(const py::array_t &rows, const py::array_t &cols, std::vector> &blocks) const override { + virtual void build_dense_blocks(const std::vector> &rows, const std::vector> &cols, std::vector> &blocks) const override { PYBIND11_OVERRIDE_PURE( void, /* Return type */ VirtualDenseBlocksGeneratorPython, /* Parent class */ @@ -49,14 +57,14 @@ class PyVirtualDenseBlocksGenerator : public VirtualDenseBlocksGeneratorPython +template > void declare_custom_VirtualDenseBlocksGenerator(py::module &m, const std::string &className) { // using BaseClass = VirtualDenseBlocksGenerator; // py::class_>(m, base_class_name.c_str()); using Class = VirtualDenseBlocksGeneratorPython; py::class_, PyVirtualDenseBlocksGenerator> py_class(m, className.c_str()); - py_class.def(py::init<>()); + py_class.def(py::init &, const Cluster &>()); py_class.def("build_dense_blocks", &Class::build_dense_blocks); } diff --git a/tests/conftest.py b/tests/conftest.py index 1ca161d..db9e62b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -141,9 +141,10 @@ def generator(geometry, cluster): params=[True, False], ids=["custom_dense_block_generator", "auto_dense_block_generator"], ) -def dense_blocks_generator(request, generator): +def dense_blocks_generator(request, generator, cluster): + [target_cluster, source_cluster] = cluster if request.param: - return CustomDenseBlocksGenerator(generator) + return CustomDenseBlocksGenerator(generator, target_cluster, source_cluster) else: return None @@ -206,7 +207,7 @@ def custom_distributed_operator( [target_cluster, source_cluster] = cluster if local_operator is not None: - custom_local_approximation = Htool.CustomApproximationBuilder( + distributed_operator_holder = Htool.CustomApproximationBuilder( target_cluster, source_cluster, symmetry, @@ -214,7 +215,6 @@ def custom_distributed_operator( mpi4py.MPI.COMM_WORLD, local_operator, ) - distributed_operator = custom_local_approximation.distributed_operator else: hmatrix_builder = Htool.HMatrixBuilder( target_cluster, @@ -232,7 +232,7 @@ def custom_distributed_operator( if low_rank_approximation is not None: hmatrix_builder.set_low_rank_generator(low_rank_approximation) - distributed_operator = Htool.DistributedOperatorFromHMatrix( + distributed_operator_holder = Htool.DistributedOperatorFromHMatrix( generator, target_cluster, source_cluster, @@ -240,7 +240,7 @@ def custom_distributed_operator( mpi4py.MPI.COMM_WORLD, ) - return distributed_operator + return distributed_operator_holder @pytest.fixture() diff --git a/tests/test_distributed_operator.py b/tests/test_distributed_operator.py index 28b11a9..a8690f4 100644 --- a/tests/test_distributed_operator.py +++ b/tests/test_distributed_operator.py @@ -21,10 +21,14 @@ (400, 400, "S", "U", True, False, False, False, False), (400, 400, "N", "N", True, False, False, False, False), (400, 200, "N", "N", True, False, False, False, False), - (400, 400, "S", "L", False, True, True, True, False), - (400, 400, "S", "U", False, True, True, True, False), - (400, 400, "N", "N", False, True, True, True, False), - (400, 200, "N", "N", False, True, True, True, False), + (400, 400, "S", "L", False, True, True, False, False), + (400, 400, "S", "U", False, True, True, False, False), + (400, 400, "N", "N", False, True, True, False, False), + (400, 200, "N", "N", False, True, True, False, False), + (400, 400, "S", "L", False, False, False, True, False), + (400, 400, "S", "U", False, False, False, True, False), + (400, 400, "N", "N", False, False, False, True, False), + (400, 200, "N", "N", False, False, False, True, False), (400, 200, "N", "N", True, False, False, False, True), ], indirect=["low_rank_approximation", "dense_blocks_generator", "local_operator"], @@ -58,9 +62,11 @@ def test_distributed_operator( fig = plt.figure() ax1 = fig.add_subplot(1, 1, 1) Htool.plot(ax1, local_hmatrix) + plt.close(fig) else: - distributed_operator = custom_distributed_operator + distributed_operator_holder = custom_distributed_operator + distributed_operator = distributed_operator_holder.distributed_operator # Test matrix vector product np.random.seed(0)