fix dense block generator

htool-ddm · Aug 9, 2024 · 213a00d · 213a00d
1 parent 948d7ed
commit 213a00d
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 61 deletions.
diff --git 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
@@ -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
@@ -7,27 +7,35 @@ namespace py = pybind11;
 using namespace pybind11::literals;
 using namespace htool;
 
-template <typename CoefficientPrecision>
+template <typename CoefficientPrecision, typename CoordinatePrecision = underlying_type<CoefficientPrecision>>
 class VirtualDenseBlocksGeneratorPython : public VirtualDenseBlocksGenerator<CoefficientPrecision> {
 
+    const Cluster<CoordinatePrecision> &m_target_cluster;
+    const Cluster<CoordinatePrecision> &m_source_cluster;
+
   public:
     using VirtualDenseBlocksGenerator<CoefficientPrecision>::VirtualDenseBlocksGenerator;
 
-    void copy_dense_blocks(const std::vector<int> &M, const std::vector<int> &N, const std::vector<int> &rows_offsets, const std::vector<int> &cols_offsets, std::vector<CoefficientPrecision *> &ptr) const override {
+    VirtualDenseBlocksGeneratorPython(const Cluster<CoordinatePrecision> &target_cluster, const Cluster<CoordinatePrecision> &source_cluster) : VirtualDenseBlocksGenerator<CoefficientPrecision>(), m_target_cluster(target_cluster), m_source_cluster(source_cluster) {}
 
-        int nb_blocks = M.size();
-        py::array_t<int> rows_offsets_np(rows_offsets.size(), rows_offsets.data(), py::capsule(rows_offsets.data()));
-        py::array_t<int> cols_offsets_np(cols_offsets.size(), cols_offsets.data(), py::capsule(cols_offsets.data()));
+    void copy_dense_blocks(const std::vector<int> &M, const std::vector<int> &N, const std::vector<int> &rows_offsets, const std::vector<int> &cols_offsets, std::vector<CoefficientPrecision *> &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<py::array_t<CoefficientPrecision, py::array::f_style>> vec_ptr;
+        std::vector<py::array_t<int, py::array::f_style>> rows_ptr;
+        std::vector<py::array_t<int, py::array::f_style>> cols_ptr;
         for (int i = 0; i < nb_blocks; i++) {
+            rows_ptr.emplace_back(std::array<long int, 1>{M[i]}, target_permutation.data() + rows_offsets[i], py::capsule(target_permutation.data() + rows_offsets[i]));
+            cols_ptr.emplace_back(std::array<long int, 1>{N[i]}, source_permutation.data() + cols_offsets[i], py::capsule(source_permutation.data() + cols_offsets[i]));
             vec_ptr.emplace_back(std::array<long int, 2>{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<int> &rows_offsets_np, const py::array_t<int> &cols_offsets_np, std::vector<py::array_t<CoefficientPrecision, py::array::f_style>> &blocks) const = 0; // LCOV_EXCL_LINE
+    virtual void build_dense_blocks(const std::vector<py::array_t<int, py::array::f_style>> &rows, const std::vector<py::array_t<int, py::array::f_style>> &cols, std::vector<py::array_t<CoefficientPrecision, py::array::f_style>> &blocks) const = 0; // LCOV_EXCL_LINE
 };
 
 template <typename CoefficientPrecision>
@@ -37,7 +45,7 @@ class PyVirtualDenseBlocksGenerator : public VirtualDenseBlocksGeneratorPython<C
     // PyVirtualGenerator(int nr0, int nc0): IMatrix<T>(nr0,nc0){}
 
     /* Trampoline (need one for each virtual function) */
-    virtual void build_dense_blocks(const py::array_t<int> &rows, const py::array_t<int> &cols, std::vector<py::array_t<CoefficientPrecision, py::array::f_style>> &blocks) const override {
+    virtual void build_dense_blocks(const std::vector<py::array_t<int, py::array::f_style>> &rows, const std::vector<py::array_t<int, py::array::f_style>> &cols, std::vector<py::array_t<CoefficientPrecision, py::array::f_style>> &blocks) const override {
         PYBIND11_OVERRIDE_PURE(
             void,                                                    /* Return type */
             VirtualDenseBlocksGeneratorPython<CoefficientPrecision>, /* Parent class */
@@ -49,14 +57,14 @@ class PyVirtualDenseBlocksGenerator : public VirtualDenseBlocksGeneratorPython<C
     }
 };
 
-template <typename CoefficientPrecision>
+template <typename CoefficientPrecision, typename CoordinatePrecision = underlying_type<CoefficientPrecision>>
 void declare_custom_VirtualDenseBlocksGenerator(py::module &m, const std::string &className) {
     // using BaseClass = VirtualDenseBlocksGenerator<CoefficientPrecision>;
     // py::class_<BaseClass, std::shared_ptr<BaseClass>>(m, base_class_name.c_str());
 
     using Class = VirtualDenseBlocksGeneratorPython<CoefficientPrecision>;
     py::class_<Class, std::shared_ptr<Class>, PyVirtualDenseBlocksGenerator<CoefficientPrecision>> py_class(m, className.c_str());
-    py_class.def(py::init<>());
+    py_class.def(py::init<const Cluster<CoordinatePrecision> &, const Cluster<CoordinatePrecision> &>());
     py_class.def("build_dense_blocks", &Class::build_dense_blocks);
 }
 

diff --git 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,15 +207,14 @@ 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,
             UPLO,
             mpi4py.MPI.COMM_WORLD,
             local_operator,
         )
-        distributed_operator = custom_local_approximation.distributed_operator
     else:
         hmatrix_builder = Htool.HMatrixBuilder(
             target_cluster,
@@ -232,15 +232,15 @@ 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,
             hmatrix_builder,
             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
@@ -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)