add distributed

Author: floatingCatty

vbcsr/atomic_data.py

@@ -104,7 +104,7 @@ def parse_param(param, name, dtype):
         )
 
     @classmethod
-    def from_distributed(
+    def from_graph_arrays(
         cls,
         n_atom,
         N_atom,
@@ -148,7 +148,7 @@ def from_distributed(
             if atomic_number_array.size != int(n_atom):
                 raise ValueError("atomic_numbers size must equal n_atom")
 
-        return super().from_distributed(
+        return super().from_graph_arrays(
             int(n_atom),
             int(N_atom),
             int(atom_offset),
@@ -165,6 +165,30 @@ def from_distributed(
             comm,
         )
 
+    @classmethod
+    def from_distributed(cls, pos, z, input_index, cell, pbc, r_max, type_norb, comm=None):
+        """Build the distributed AtomicData from a CALLER-GIVEN partition (doc 42 §4).
+
+        Each rank passes ONLY its owned atoms: ``pos`` (n_owned, 3), ``z`` (n_owned,),
+        and ``input_index`` (n_owned,) — the original input-order index of each owned
+        atom (fills ``atom_index``/``indices``). ``r_max``/``type_norb`` are per-type
+        (global, indexed by sorted-unique-Z). Global ids are assigned contiguously by
+        rank and the edge graph + ghosts are built distributed (no rank-0 gather, no
+        ParMETIS). The partition (which atom each rank owns) is the caller's choice.
+        """
+        comm = _default_comm(comm)
+        pos = np.ascontiguousarray(np.asarray(pos, dtype=np.float64).reshape(-1, 3))
+        z = np.ascontiguousarray(np.asarray(z, dtype=np.int32).reshape(-1))
+        input_index = np.ascontiguousarray(np.asarray(input_index, dtype=np.int32).reshape(-1))
+        cell = np.ascontiguousarray(np.asarray(cell, dtype=np.float64).reshape(3, 3))
+        r_max_vec = np.ascontiguousarray(np.asarray(r_max, dtype=np.float64).reshape(-1))
+        type_norb_vec = np.ascontiguousarray(np.asarray(type_norb, dtype=np.int32).reshape(-1))
+        if z.size != pos.shape[0] or input_index.size != pos.shape[0]:
+            raise ValueError("from_distributed: pos/z/input_index must share n_owned")
+        return super().from_distributed(
+            pos, z, input_index, cell, _normalize_pbc(pbc), r_max_vec, type_norb_vec, comm
+        )
+
     @classmethod
     def from_ase(cls, atoms, r_max, type_norb=1, comm=None):
         pbc = atoms.pbc

vbcsr/core/atomic/atomic_data.hpp

@@ -2,6 +2,8 @@
 #include "../dist_graph.hpp"
 #include <vector>
 #include <map>
+#include <set>
+#include <array>
 #include <algorithm>
 #include <mpi.h>
 #include <iostream>
@@ -370,7 +372,80 @@ class AtomicData {
             r_edges,
             type_norb);
     }
-    
+
+    // Distributed construction from a CALLER-GIVEN partition (doc/design/42 §4).
+    // Each rank passes ONLY its owned atoms — there is no rank-0 gather of the
+    // geometry and no ParMETIS: the partition is whatever the caller chose (e.g.
+    // a spatial slab/RCB aligned with the grid). Global ids are assigned
+    // contiguously by rank (Allgather of counts -> Exscan), so rank r owns
+    // [offset_r, offset_r + n_r). This form Allgathers positions/Z (O(N), transient)
+    // so every rank builds the edges of ITS OWNED atoms locally; the persistent
+    // storage is owned + ghost only. ``my_input_index[li]`` is the original
+    // input-order index of owned atom ``li`` (fills AtomicData.atom_index).
+    static AtomicData* from_distributed(
+        const std::vector<double>& my_pos, const std::vector<int>& my_z,
+        const std::vector<int>& my_input_index,
+        const std::vector<double>& cell, const std::vector<bool>& pbc,
+        const std::vector<double>& r_max_per_type, const std::vector<int>& type_norb_in,
+        MPI_Comm comm) {
+        int rank = 0, size = 1, initialized = 0;
+        MPI_Initialized(&initialized);
+        if (initialized) { MPI_Comm_rank(comm, &rank); MPI_Comm_size(comm, &size); }
+        const int my_n = static_cast<int>(my_z.size());
+
+        // 1. Contiguous global-id layout: rank r owns [offset_r, offset_r + n_r).
+        std::vector<int> counts(size, my_n), displs(size + 1, 0);
+        if (initialized) MPI_Allgather(&my_n, 1, MPI_INT, counts.data(), 1, MPI_INT, comm);
+        for (int i = 0; i < size; ++i) displs[i + 1] = displs[i] + counts[i];
+        const int total = displs[size];
+        const int my_offset = displs[rank];
+        if (total == 0) return new AtomicData(comm);
+
+        // 2. Allgather all positions + atomic numbers (global-id == rank order).
+        std::vector<int> pos_counts(size), pos_displs(size + 1, 0);
+        for (int i = 0; i < size; ++i) { pos_counts[i] = counts[i] * 3; pos_displs[i + 1] = pos_displs[i] + pos_counts[i]; }
+        std::vector<double> all_pos(static_cast<size_t>(total) * 3);
+        std::vector<int> all_z(total);
+        if (initialized) {
+            MPI_Allgatherv(my_pos.data(), my_n * 3, MPI_DOUBLE, all_pos.data(), pos_counts.data(), pos_displs.data(), MPI_DOUBLE, comm);
+            MPI_Allgatherv(my_z.data(), my_n, MPI_INT, all_z.data(), counts.data(), displs.data(), MPI_INT, comm);
+        } else { all_pos = my_pos; all_z = my_z; }
+
+        // 3. Global z -> type (sorted-unique z; identical on every rank).
+        std::vector<int> uz = all_z; std::sort(uz.begin(), uz.end());
+        uz.erase(std::unique(uz.begin(), uz.end()), uz.end());
+        std::map<int, int> z2t; for (size_t i = 0; i < uz.size(); ++i) z2t[uz[i]] = static_cast<int>(i);
+        std::vector<int> all_types(total); for (int g = 0; g < total; ++g) all_types[g] = z2t[all_z[g]];
+        std::vector<int> type_norb = type_norb_in;
+        if (type_norb.empty()) type_norb.assign(uz.size(), 1);
+        std::vector<int> my_types(my_n); for (int li = 0; li < my_n; ++li) my_types[li] = all_types[my_offset + li];
+
+        // 4. Edges of MY owned atoms (NeighborList over all positions; same cutoff
+        //    + distance test as process_input_rank0, so the graph is identical to
+        //    from_points for the same geometry). r_edges: 5 ints {gi, gj, rx, ry, rz}.
+        double max_r = 0; for (double r : r_max_per_type) max_r = std::max(max_r, r);
+        NeighborList nl; nl.build(all_pos, cell, pbc, max_r * 2.0);
+        std::vector<int> r_edges;
+        for (int li = 0; li < my_n; ++li) {
+            const int gi = my_offset + li, ti = all_types[gi];
+            for (const auto& nb : nl.get_neighbors(gi)) {
+                const int gj = nb.index, tj = all_types[gj];
+                const double rc = r_max_per_type[ti] + r_max_per_type[tj];
+                const double dx = all_pos[3 * gj]     - all_pos[3 * gi]     + nb.rx * cell[0] + nb.ry * cell[3] + nb.rz * cell[6];
+                const double dy = all_pos[3 * gj + 1] - all_pos[3 * gi + 1] + nb.rx * cell[1] + nb.ry * cell[4] + nb.rz * cell[7];
+                const double dz = all_pos[3 * gj + 2] - all_pos[3 * gi + 2] + nb.rx * cell[2] + nb.ry * cell[5] + nb.rz * cell[8];
+                if (std::sqrt(dx * dx + dy * dy + dz * dz) > rc + 1e-9) continue;
+                r_edges.push_back(gi); r_edges.push_back(gj);
+                r_edges.push_back(nb.rx); r_edges.push_back(nb.ry); r_edges.push_back(nb.rz);
+            }
+        }
+
+        // 5. Build the distributed object (owned atoms + owned-src edges; the ctor
+        //    derives ghosts from the edge dsts). Same final assembly as from_points.
+        return construct_final_object(comm, rank, size, cell, pbc, my_n,
+            my_input_index, my_z, my_types, my_pos, r_edges, type_norb);
+    }
+
     static AtomicData* from_file(const std::string& filename, const std::vector<double>& r_max_per_type, std::vector<int> type_norb, MPI_Comm comm, const std::string& format="") {
         int rank;
         int initialized = 0;
@@ -749,10 +824,144 @@ class AtomicData {
         }
 
         new_graph->construct_distributed(owned_indices, my_block_sizes, matrix_adj);
-        
+
         return new_graph;
     }
 
+    // Build a full ``AtomicData`` whose neighbour graph IS the 3-centre (graph3b) structure,
+    // with the SAME atom distribution as ``this`` — the proper-AtomicData sibling of
+    // ``get_graph3b`` (doc/design/41). Unlike ``get_graph3b`` (which returns a shift-less
+    // ``DistGraph`` and so forces the union-graph ``ImageContainer`` ctor + a hand-built
+    // explicit shift set), this tracks each 3-centre block's integer lattice shift, so the
+    // normal ``ImageContainer(AtomicData*)`` ctor builds the per-R image graphs directly.
+    //
+    // For an owned centre i with reduced neighbours (j, R_j) and (k, R_k) (within the
+    // ``r_max_left[itype] + r_max_right[jtype]`` cut, same as get_graph3b), the 3-centre block
+    // (gid_j -> gid_k) lives at shift R_k - R_j and is owned by gid_j's row owner; the routing
+    // is the same one-Alltoallv idiom as get_graph3b, carrying 5 ints (gj, gk, dRx, dRy, dRz)
+    // per block. The self-onsite (j==k, R=0) is dropped — the ImageContainer ctor re-adds the
+    // R=0 diagonal. The resulting (gi, gj, R) edge set equals the graph3b pair set the V_nl /
+    // force / DM consumers enumerate, so an ImageContainer on this AtomicData has exactly their
+    // blocks (and ``build_dmr_images`` can read its pairs straight from the graph).
+    AtomicData* get_atomicdata3b(const std::vector<double>& r_max_left,
+                                 const std::vector<double>& r_max_right) {
+        int initialized = 0;
+        MPI_Initialized(&initialized);
+
+        // 1. Reduced neighbour list WITH integer shift (do NOT collapse periodic images).
+        std::vector<std::vector<std::pair<int, std::array<int, 3>>>> rnbr(n_atom);
+        for (int i = 0; i < n_atom; ++i) {
+            const int itype = atom_type[i];
+            for (int edge_idx : iconn[i]) {
+                const int j = edges[edge_idx].dst;
+                double rx, ry, rz;
+                get_edge_vec(edge_idx, &rx, &ry, &rz);
+                const double r = std::sqrt(rx * rx + ry * ry + rz * rz);
+                const int jtype = atom_type[j];
+                if (r <= r_max_left[itype] + r_max_right[jtype] + 1e-9) {
+                    int sx, sy, sz;
+                    get_edge_shift_vec(edge_idx, &sx, &sy, &sz);
+                    rnbr[i].push_back({get_gid(j), {sx, sy, sz}});
+                }
+            }
+        }
+
+        // 2. Each ordered neighbour pair (j,Rj),(k,Rk) of i -> block (gid_j -> gid_k) at Rk-Rj,
+        //    routed to gid_j's owner (5 ints each).
+        std::map<int, std::vector<std::array<int, 5>>> send_map;
+        for (int i = 0; i < n_atom; ++i) {
+            const auto& nb = rnbr[i];
+            for (const auto& a : nb) {
+                const int owner_j = graph->find_owner(a.first);
+                for (const auto& b : nb) {
+                    send_map[owner_j].push_back({a.first, b.first,
+                        b.second[0] - a.second[0], b.second[1] - a.second[1], b.second[2] - a.second[2]});
+                }
+            }
+        }
+        for (auto& kv : send_map) {
+            std::sort(kv.second.begin(), kv.second.end());
+            kv.second.erase(std::unique(kv.second.begin(), kv.second.end()), kv.second.end());
+        }
+
+        std::vector<int> send_counts(size, 0);
+        for (auto& kv : send_map) send_counts[kv.first] = static_cast<int>(kv.second.size()) * 5;
+        std::vector<int> recv_counts(size);
+        if (initialized) {
+            MPI_Alltoall(send_counts.data(), 1, MPI_INT, recv_counts.data(), 1, MPI_INT, comm);
+        } else {
+            recv_counts = send_counts;
+        }
+        std::vector<int> sdispls(size + 1, 0), rdispls(size + 1, 0);
+        for (int i = 0; i < size; ++i) {
+            sdispls[i + 1] = sdispls[i] + send_counts[i];
+            rdispls[i + 1] = rdispls[i] + recv_counts[i];
+        }
+        std::vector<int> send_buf(sdispls[size]);
+        for (auto& kv : send_map) {
+            int off = sdispls[kv.first];
+            for (const auto& e : kv.second) for (int t = 0; t < 5; ++t) send_buf[off++] = e[t];
+        }
+        std::vector<int> recv_buf(rdispls[size]);
+        if (initialized) {
+            MPI_Alltoallv(send_buf.data(), send_counts.data(), sdispls.data(), MPI_INT,
+                          recv_buf.data(), recv_counts.data(), rdispls.data(), MPI_INT, comm);
+        } else {
+            recv_buf = send_buf;
+        }
+
+        // 3. Dedup received 3-centre edges (gid_j owned here); drop the self-onsite.
+        std::set<std::array<int, 5>> edge_set;
+        for (size_t p = 0; p + 4 < recv_buf.size(); p += 5) {
+            if (recv_buf[p] == recv_buf[p + 1] &&
+                recv_buf[p + 2] == 0 && recv_buf[p + 3] == 0 && recv_buf[p + 4] == 0) continue;
+            edge_set.insert({recv_buf[p], recv_buf[p + 1], recv_buf[p + 2], recv_buf[p + 3], recv_buf[p + 4]});
+        }
+
+        // graph3b ⊇ the 2-body S/T graph: add this rank's owned 2-body edges (i -> j, R_j) WITH
+        // shift (owned locally, no routing). The 3-centre enumeration alone covers only pairs
+        // sharing a common centre, so the pure 2-body pairs must be added explicitly — mirrors
+        // get_graph3b's "add two-body and onsite edges". The R=0 onsite is re-added by the
+        // ImageContainer ctor, so it is skipped here.
+        for (int i = 0; i < n_atom; ++i) {
+            const int gi = get_gid(i);
+            for (int edge_idx : iconn[i]) {
+                const int j = edges[edge_idx].dst;
+                int sx, sy, sz;
+                get_edge_shift_vec(edge_idx, &sx, &sy, &sz);
+                if (gi == get_gid(j) && sx == 0 && sy == 0 && sz == 0) continue;
+                edge_set.insert({gi, get_gid(j), sx, sy, sz});
+            }
+        }
+
+        std::vector<int> edge_index, edge_shift;
+        edge_index.reserve(edge_set.size() * 2);
+        edge_shift.reserve(edge_set.size() * 3);
+        for (const auto& e : edge_set) {
+            edge_index.push_back(e[0]); edge_index.push_back(e[1]);
+            edge_shift.push_back(e[2]); edge_shift.push_back(e[3]); edge_shift.push_back(e[4]);
+        }
+        int n_edge3 = static_cast<int>(edge_set.size());
+        int N_edge3 = n_edge3;
+        if (initialized) MPI_Allreduce(&n_edge3, &N_edge3, 1, MPI_INT, MPI_SUM, comm);
+
+        // 4. Owned per-atom arrays copied from this AtomicData.
+        std::vector<int> at_index(n_atom), at_type(n_atom), at_num(n_atom);
+        std::vector<double> pos(3 * n_atom);
+        for (int i = 0; i < n_atom; ++i) {
+            at_index[i] = atom_index[i];
+            at_type[i] = atom_type[i];
+            at_num[i] = atomic_numbers[i];
+            pos[3 * i] = x[i]; pos[3 * i + 1] = y[i]; pos[3 * i + 2] = z[i];
+        }
+
+        return new AtomicData(
+            static_cast<size_t>(n_atom), static_cast<size_t>(N_atom), static_cast<size_t>(atom_offset),
+            static_cast<size_t>(n_edge3), static_cast<size_t>(N_edge3),
+            at_index.data(), at_type.data(), edge_index.data(), type_norb.data(), edge_shift.data(),
+            cell.data(), pos.data(), at_num.data(), comm, pbc);
+    }
+
 private:
     static std::vector<bool> infer_pbc_from_edge_shifts(size_t n_edge, const int *edge_shift_vec_in, int initialized, MPI_Comm comm) {
         int local_flags[3] = {0, 0, 0};

vbcsr/core/atomic/image_container.hpp

@@ -45,30 +45,20 @@ class ImageContainer {
     // True when this container owns ``base_graph`` (the union-graph ctor below);
     // the normal ctor borrows ``data->graph`` and leaves this false.
     bool owns_base_graph = false;
+    // True when this container was built with the owning ``ImageContainer(AtomicData*, true)``
+    // ctor and must delete ``atom_data`` (and, through it, its graph) in the destructor.
+    bool owns_atom_data = false;
 
 public:
-    ImageContainer(AtomicData* data) : atom_data(data), base_graph(data->graph) {
+    // ``own_atom_data`` lets the container take ownership of a freshly-built AtomicData (e.g.
+    // ``get_atomicdata3b``, doc/design/41) and delete it in the destructor — so a graph3b
+    // operator/weight container can be built with this NORMAL ctor (which reads the per-edge
+    // shifts to build the per-R image graphs) instead of the union-graph ctor below.
+    ImageContainer(AtomicData* data, bool own_atom_data = false)
+        : atom_data(data), base_graph(data->graph), owns_atom_data(own_atom_data) {
         build_image_graphs();
     }
 
-    // Union-graph ctor: every image (one per shift in ``shifts``, which MUST be the
-    // globally-consistent shift set across the comm) lives on the single provided
-    // ``union_graph`` rather than per-R exact graphs derived from the 2-body
-    // ``iconn``. This is for operators whose (i, j) sparsity is wider than the
-    // AtomicData edges — notably V_nl, whose pairs share a common projector and so
-    // span the 3-body ``get_graph3b`` graph. ``add_block`` routes each (i, j) block
-    // to its row owner exactly like ``VBCSR(graph3b).add_block``; ``sample_k``
-    // accumulates onto ``base_graph = union_graph``. If ``own_union`` the container
-    // deletes ``union_graph`` in its destructor. The per-R ``BlockSpMat`` borrow the
-    // graph (non-owning), so there is no double free.
-    ImageContainer(AtomicData* data, DistGraph* union_graph,
-                   const std::vector<std::vector<int>>& shifts, bool own_union = false)
-        : atom_data(data), base_graph(union_graph), owns_base_graph(own_union) {
-        for (const auto& R : shifts) {
-            image_blocks[R] = new BlockSpMat<T, Kernel>(union_graph);
-        }
-    }
-
     ~ImageContainer() {
         for (auto& kv : image_blocks) {
             delete kv.second;
@@ -79,6 +69,9 @@ class ImageContainer {
         if (owns_base_graph) {
             delete base_graph;
         }
+        if (owns_atom_data) {
+            delete atom_data;  // deletes its own_graph too (AtomicData dtor)
+        }
     }
 
     void build_image_graphs() {
@@ -233,6 +226,25 @@ class ImageContainer {
         }
     }
 
+    // In-place per-image axpy: ``this += alpha * other``, matching images by their R shift. Used
+    // to fuse the 2-body kinetic T into the graph3b V_nl to form the combined static Hamiltonian
+    // H_static = T + V_nl (doc/design/41 §3), so only {S, H_static} is sent down (not 3 images)
+    // and V_nl is never rebuilt on the pool graph. ``this`` MUST be a superset graph — every R in
+    // ``other`` present here, and (per-R) every (row,col) of ``other`` present in this image — the
+    // graph3b ⊇ 2-body case. Both sides must share the owned-row partition + block sizes; the
+    // cross-graph ``BlockSpMat::axpby`` maps blocks by GLOBAL (row,col) index.
+    void axpy_into(const ImageContainer& other, T alpha) {
+        for (const auto& kv : other.image_blocks) {
+            auto it = image_blocks.find(kv.first);
+            if (it == image_blocks.end()) {
+                throw std::runtime_error(
+                    "ImageContainer::axpy_into: a shift R of `other` is absent in this container "
+                    "(this must be a superset graph, e.g. graph3b absorbing the 2-body operator)");
+            }
+            it->second->axpy(alpha, *kv.second);
+        }
+    }
+
     // Read one image block (R, global row, global col) on the owner of the row;
     // returns empty if this rank does not hold it. RowMajor by default.
     std::vector<T> get_block(const std::vector<int>& R, int g_row, int g_col,
@@ -361,6 +373,12 @@ class ImageContainer {
         }
     }
 
+    // NOTE: the former ``gather_into`` (request-list gather that re-``assemble``d into a target
+    // DistGraph) was removed in doc/design/41 §2.3 — its ``assemble`` re-routes the just-fetched
+    // blocks by the target partition (a no-op only for a serial target, and an un-gather for any
+    // real partition). The operator-walk gather is now the keyed ``LocalBlockView`` produced by
+    // ``rescupp::lcao_grid::GatherPlan::gather_local`` (alias local, cache remote, no re-route).
+
     // Accumulate all image blocks onto a compatible reference graph using
     // a per-image weight and an optional per-block correction factor.
     template <typename ResultT, typename ResultKernel = DefaultKernel<ResultT>, typename ImageWeightFn, typename BlockWeightFn>