fixup

Author: PierreMarchand20

example/define_custom_local_operator.py

@@ -2,30 +2,32 @@
 import numpy as np
 
 
-class CustomLocalOperator(Htool.VirtualLocalOperator):
+class CustomLocalOperator(Htool.LocalOperator):
     def __init__(
         self,
         generator: Htool.VirtualGenerator,
-        target_cluster: Htool.Cluster,
-        global_source_cluster: Htool.Cluster,
-        local_source_offset: int,
-        local_source_size: int,
+        target_local_renumbering: Htool.LocalRenumbering,
+        source_local_renumbering: Htool.LocalRenumbering,
+        target_use_permutation_to_mvprod: bool = False,
+        source_use_permutation_to_mvprod: bool = False,
     ) -> None:
         super().__init__(
-            target_cluster,
-            global_source_cluster,
-            local_source_offset,
-            local_source_size,
+            target_local_renumbering,
+            source_local_renumbering,
+            target_use_permutation_to_mvprod,
+            source_use_permutation_to_mvprod,
+        )
+        self.data = np.zeros(
+            (target_local_renumbering.size, source_local_renumbering.size)
         )
-        target_offset = target_cluster.get_offset()
-        target_size = target_cluster.get_size()
-        self.data = np.zeros((target_size, local_source_size))
         generator.build_submatrix(
-            target_cluster.get_permutation()[
-                target_offset : target_offset + target_size
+            target_local_renumbering.permutation[
+                target_local_renumbering.offset : target_local_renumbering.offset
+                + target_local_renumbering.size
             ],
-            global_source_cluster.get_permutation()[
-                local_source_offset : local_source_offset + local_source_size
+            source_local_renumbering.permutation[
+                source_local_renumbering.offset : source_local_renumbering.offset
+                + source_local_renumbering.size
             ],
             self.data,
         )
@@ -36,60 +38,20 @@ def add_vector_product(
         # Beware, inplace operation needed for output to keep the underlying data
         output *= beta
         if trans == "N":
-            output += alpha * self.data.dot(
-                input[
-                    self.local_source_offset : self.local_source_offset
-                    + self.local_source_size
-                ]
-            )
+            output += alpha * self.data.dot(input)
         elif trans == "T":
-            output += alpha * np.transpose(self.data).dot(
-                input[
-                    self.local_source_offset : self.local_source_offset
-                    + self.local_source_size
-                ]
-            )
+            output += alpha * np.transpose(self.data).dot(input)
         elif trans == "C":
-            output += alpha * np.vdot(
-                np.transpose(self.data),
-                input[
-                    self.local_source_offset : self.local_source_offset
-                    + self.local_source_size
-                ],
-            )
+            output += alpha * np.vdot(np.transpose(self.data), input)
 
     def add_matrix_product_row_major(
         self, trans, alpha, input: np.array, beta, output: np.array
     ) -> None:
         output *= beta
         if trans == "N":
-            output += (
-                alpha
-                * self.data
-                @ input[
-                    self.local_source_offset : self.local_source_offset
-                    + self.local_source_size,
-                    :,
-                ]
-            )
+            output += alpha * self.data @ input
         elif trans == "T":
-            output += (
-                alpha
-                * np.transpose(self.data)
-                @ input[
-                    self.local_source_offset : self.local_source_offset
-                    + self.local_source_size,
-                    :,
-                ]
-            )
+            output += alpha * np.transpose(self.data) @ input
         elif trans == "C":
-            output += (
-                alpha
-                * np.matrix.H(self.data)
-                @ input[
-                    self.local_source_offset : self.local_source_offset
-                    + self.local_source_size,
-                    :,
-                ]
-            )
+            output += alpha * np.matrix.H(self.data) @ input
         output = np.asfortranarray(output)

example/use_local_hmatrix_compression.py

@@ -76,10 +76,10 @@
 if local_source_cluster.get_offset() > 0:
     local_operator_1 = CustomLocalOperator(
         generator,
-        local_target_cluster,
-        source_cluster,
-        0,
-        local_source_cluster.get_offset(),
+        Htool.LocalRenumbering(local_target_cluster),
+        Htool.LocalRenumbering(
+            0, local_source_cluster.get_offset(), source_cluster.get_permutation()
+        ),
     )
 
 local_operator_2 = None
@@ -91,12 +91,14 @@
 ):
     local_operator_2 = CustomLocalOperator(
         generator,
-        local_target_cluster,
-        source_cluster,
-        local_source_cluster.get_size() + local_source_cluster.get_offset(),
-        source_cluster.get_size()
-        - local_source_cluster.get_size()
-        - local_source_cluster.get_offset(),
+        Htool.LocalRenumbering(local_target_cluster),
+        Htool.LocalRenumbering(
+            local_source_cluster.get_size() + local_source_cluster.get_offset(),
+            source_cluster.get_size()
+            - local_source_cluster.get_size()
+            - local_source_cluster.get_offset(),
+            source_cluster.get_permutation(),
+        ),
     )
 
 if local_operator_1:

lib/htool

@@ -4,25 +4,25 @@
 #include <htool/distributed_operator/implementations/local_operators/local_operator.hpp>
 #include <pybind11/pybind11.h>
 
-template <typename CoefficientPrecision, typename CoordinatePrecision = CoefficientPrecision>
-class LocalOperatorPython : public htool::LocalOperator<CoefficientPrecision, CoordinatePrecision> {
+template <typename CoefficientPrecision>
+class LocalOperatorPython : public htool::LocalOperator<CoefficientPrecision> {
   public:
-    using htool::LocalOperator<CoefficientPrecision, CoordinatePrecision>::LocalOperator;
+    using htool::LocalOperator<CoefficientPrecision>::LocalOperator;
 
-    LocalOperatorPython(const Cluster<CoordinatePrecision> &cluster_tree_target, const Cluster<CoordinatePrecision> &cluster_tree_source, bool target_use_permutation_to_mvprod = false, bool source_use_permutation_to_mvprod = false) : LocalOperator<CoefficientPrecision, CoordinatePrecision>(cluster_tree_target, cluster_tree_source, target_use_permutation_to_mvprod, source_use_permutation_to_mvprod) {}
+    LocalOperatorPython(LocalRenumbering target_local_renumbering, LocalRenumbering source_local_renumbering, bool target_use_permutation_to_mvprod = false, bool source_use_permutation_to_mvprod = false) : LocalOperator<CoefficientPrecision>(target_local_renumbering, source_local_renumbering, target_use_permutation_to_mvprod, source_use_permutation_to_mvprod) {}
 
     void local_add_vector_product(char trans, CoefficientPrecision alpha, const CoefficientPrecision *in, CoefficientPrecision beta, CoefficientPrecision *out) const override {
 
-        py::array_t<CoefficientPrecision> input(std::array<long int, 1>{this->m_source_cluster.get_size()}, in, py::capsule(in));
-        py::array_t<CoefficientPrecision> output(std::array<long int, 1>{this->m_target_cluster.get_size()}, out, py::capsule(out));
+        py::array_t<CoefficientPrecision> input(std::array<long int, 1>{this->m_local_source_renumbering.get_size()}, in, py::capsule(in));
+        py::array_t<CoefficientPrecision> output(std::array<long int, 1>{this->m_local_target_renumbering.get_size()}, out, py::capsule(out));
 
         add_vector_product(trans, alpha, input, beta, output);
     }
 
     void local_add_matrix_product_row_major(char trans, CoefficientPrecision alpha, const CoefficientPrecision *in, CoefficientPrecision beta, CoefficientPrecision *out, int mu) const override {
 
-        py::array_t<CoefficientPrecision, py::array::c_style> input(std::array<long int, 2>{this->m_source_cluster.get_size(), mu}, in, py::capsule(in));
-        py::array_t<CoefficientPrecision, py::array::c_style> output(std::array<long int, 2>{this->m_target_cluster.get_size(), mu}, out, py::capsule(out));
+        py::array_t<CoefficientPrecision, py::array::c_style> input(std::array<long int, 2>{this->m_local_source_renumbering.get_size(), mu}, in, py::capsule(in));
+        py::array_t<CoefficientPrecision, py::array::c_style> output(std::array<long int, 2>{this->m_local_target_renumbering.get_size(), mu}, out, py::capsule(out));
 
         add_matrix_product_row_major(trans, alpha, input, beta, output);
     }
@@ -34,10 +34,10 @@ class LocalOperatorPython : public htool::LocalOperator<CoefficientPrecision, Co
     virtual void add_matrix_product_row_major(char trans, CoefficientPrecision alpha, const py::array_t<CoefficientPrecision, py::array::c_style> &in, CoefficientPrecision beta, py::array_t<CoefficientPrecision, py::array::c_style> &out) const = 0; // LCOV_EXCL_LINE
 };
 
-template <typename CoefficientPrecision, typename CoordinatePrecision>
-class PyLocalOperator : public LocalOperatorPython<CoefficientPrecision, CoordinatePrecision> {
+template <typename CoefficientPrecision>
+class PyLocalOperator : public LocalOperatorPython<CoefficientPrecision> {
   public:
-    using LocalOperatorPython<CoefficientPrecision, CoordinatePrecision>::LocalOperatorPython;
+    using LocalOperatorPython<CoefficientPrecision>::LocalOperatorPython;
 
     /* Trampoline (need one for each virtual function) */
     virtual void add_vector_product(char trans, CoefficientPrecision alpha, const py::array_t<CoefficientPrecision> &in, CoefficientPrecision beta, py::array_t<CoefficientPrecision> &out) const override {
@@ -66,17 +66,17 @@ class PyLocalOperator : public LocalOperatorPython<CoefficientPrecision, Coordin
     }
 };
 
-template <typename CoefficientPrecision, typename CoordinatePrecision>
+template <typename CoefficientPrecision>
 void declare_local_operator(py::module &m, const std::string &class_name) {
     using VirtualClass = htool::VirtualLocalOperator<CoefficientPrecision>;
     py::class_<VirtualClass>(m, ("Virtual" + class_name).c_str());
 
-    using BaseClass = LocalOperator<CoefficientPrecision, CoordinatePrecision>;
+    using BaseClass = LocalOperator<CoefficientPrecision>;
     py::class_<BaseClass, VirtualClass> py_base_class(m, ("Base" + class_name).c_str());
 
-    using Class = LocalOperatorPython<CoefficientPrecision, CoordinatePrecision>;
-    py::class_<Class, PyLocalOperator<CoefficientPrecision, CoordinatePrecision>, BaseClass> py_class(m, class_name.c_str());
-    py_class.def(py::init<const Cluster<CoordinatePrecision> &, const Cluster<CoordinatePrecision> &, bool, bool>());
+    using Class = LocalOperatorPython<CoefficientPrecision>;
+    py::class_<Class, PyLocalOperator<CoefficientPrecision>, BaseClass> py_class(m, class_name.c_str());
+    py_class.def(py::init<LocalRenumbering, LocalRenumbering, bool, bool>());
     py_class.def("add_vector_product", &Class::add_vector_product, py::arg("trans"), py::arg("alpha"), py::arg("in").noconvert(true), py::arg("beta"), py::arg("out").noconvert(true));
     py_class.def("add_matrix_product_row_major", &Class::add_matrix_product_row_major);
 }

src/htool/local_operator/local_operator.hpp

@@ -0,0 +1,21 @@
+#ifndef HTOOL_VIRTUAL_LOCAL_RENUMBERING_CPP
+#define HTOOL_VIRTUAL_LOCAL_RENUMBERING_CPP
+
+#include <htool/distributed_operator/local_renumbering.hpp>
+#include <pybind11/pybind11.h>
+
+template <typename CoordinatePrecision>
+void declare_local_renumbering(py::module &m, const std::string &className) {
+
+    using Class = LocalRenumbering;
+    py::class_<Class> py_class(m, className.c_str());
+    py_class.def(py::init([](int offset, int size, py::array_t<int> permutation) {
+        return std::unique_ptr<Class>(new Class(offset, size, permutation.size(), permutation.data()));
+    }));
+    py_class.def(py::init<const Cluster<CoordinatePrecision> &>());
+    py_class.def_property_readonly("offset", &Class::get_offset);
+    py_class.def_property_readonly("size", &Class::get_size);
+    py_class.def_property_readonly("permutation", [](const Class &self) { return py::array_t<int>(std::array<long int, 1>{self.get_global_size()}, self.get_permutation(), py::capsule(self.get_permutation())); });
+}
+
+#endif

src/htool/local_operator/local_renumbering.hpp

@@ -16,6 +16,7 @@
 #include "hmatrix/lrmat.hpp"
 
 #include "local_operator/local_operator.hpp"
+#include "local_operator/local_renumbering.hpp"
 #include "local_operator/virtual_local_operator.hpp"
 
 #include "distributed_operator/distributed_operator.hpp"
@@ -58,7 +59,8 @@ PYBIND11_MODULE(Htool, m) {
     declare_custom_VirtualLowRankGenerator<double>(m, "VirtualLowRankGenerator");
     declare_custom_VirtualDenseBlocksGenerator<double>(m, "VirtualDenseBlocksGenerator");
 
-    declare_virtual_local_operator<double>(m, "VirtualLocalOperator", "ILocalOperator");
+    declare_local_renumbering<double>(m, "LocalRenumbering");
+    declare_local_operator<double>(m, "LocalOperator");
     // declare_local_operator<double, double>(m, "LocalOperator");
     declare_distributed_operator<double>(m, "DistributedOperator");
     declare_distributed_operator_utility<double, double>(m);

src/htool/main.cpp

@@ -151,7 +151,7 @@ def cluster(geometry, symmetry):
             mpi4py.MPI.COMM_WORLD.size,
             partition=None,
             radii=None,
-            weights=None
+            weights=None,
         )
 
     if target_partition is not None:
@@ -161,7 +161,7 @@ def cluster(geometry, symmetry):
             mpi4py.MPI.COMM_WORLD.size,
             partition=target_partition,
             radii=None,
-            weights=None
+            weights=None,
         )
     else:
         target_cluster: Htool.Cluster = cluster_builder.create_cluster_tree(
@@ -170,7 +170,7 @@ def cluster(geometry, symmetry):
             mpi4py.MPI.COMM_WORLD.size,
             partition=None,
             radii=None,
-            weights=None
+            weights=None,
         )
 
     if symmetry == "S" or symmetry == "H":
@@ -205,11 +205,10 @@ def local_operator(request, generator, cluster, geometry):
         [target_cluster, source_cluster] = cluster
         return CustomLocalOperator(
             generator,
-            target_cluster.get_cluster_on_partition(mpi4py.MPI.COMM_WORLD.rank),
-            source_cluster,
-            source_cluster.get_offset(),
-            source_cluster.get_size(),
-
+            Htool.LocalRenumbering(
+                target_cluster.get_cluster_on_partition(mpi4py.MPI.COMM_WORLD.rank)
+            ),
+            Htool.LocalRenumbering(source_cluster),
         )
     else:
         return None