Skip to content

Commit 6df5c6d

Browse files
jinzx10jinzx10
authored
Grid batching by adapted cut-plane (#5599)
* initial commit * continue... * crude test... * continue... * complete! * more test * small fix
1 parent cab5894 commit 6df5c6d

File tree

5 files changed

+433
-1
lines changed

5 files changed

+433
-1
lines changed

‎source/module_base/blas_connector.h‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,20 @@ extern "C"
3939
double dnrm2_( const int *n, const double *X, const int *incX );
4040
double dznrm2_( const int *n, const std::complex<double> *X, const int *incX );
4141

42+
// symmetric rank-k update
43+
void dsyrk_(
44+
const char* uplo,
45+
const char* trans,
46+
const int* n,
47+
const int* k,
48+
const double* alpha,
49+
const double* a,
50+
const int* lda,
51+
const double* beta,
52+
double* c,
53+
const int* ldc
54+
);
55+
4256
// level 2: matrix-std::vector operations, O(n^2) data and O(n^2) work.
4357
void sgemv_(const char*const transa, const int*const m, const int*const n,
4458
const float*const alpha, const float*const a, const int*const lda, const float*const x, const int*const incx,
@@ -267,4 +281,4 @@ void zgemv_i(const char *trans,
267281
*/
268282

269283
#endif // GATHER_INFO
270-
#endif // BLAS_CONNECTOR_H
284+
#endif // BLAS_CONNECTOR_H

‎source/module_base/grid/batch.cpp‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
#include "module_base/grid/batch.h"
2+
3+
#include <algorithm>
4+
#include <cassert>
5+
#include <iterator>
6+
7+
#include "module_base/blas_connector.h"
8+
#include "module_base/lapack_connector.h"
9+
10+
namespace {
11+
12+
/**
13+
* @brief Divide a set of points into two subsets by the "MaxMin" algorithm.
14+
*
15+
* This function divides a given set of grid points by a cut plane
16+
* {x|n^T*(x-c) = 0} where the normal vector n and the point c are
17+
* determined by the "MaxMin" problem:
18+
*
19+
* max min sum_{i=1}^{m} [n^T * (r[idx[i]] - c)]^2
20+
* n c
21+
*
22+
* here r[j] = (grid[3*j], grid[3*j+1], grid[3*j+2]) is the position of
23+
* the j-th point.
24+
*
25+
* It can be shown that the optimal c is the centroid of the points, and
26+
* the optimal n is the eigenvector corresponding to the largest eigenvalue
27+
* of the matrix R*R^T, where the i-th column of R is r[idx[i]] - c.
28+
*
29+
* @param[in] grid Coordinates of all grid points.
30+
* grid[3*j], grid[3*j+1], grid[3*j+2] are the
31+
* x, y, z coordinates of the j-th point.
32+
* @param[in,out] idx Indices of the selected points within grid.
33+
* On return, idx will be rearranged such that
34+
* points belonging to the same subset have their
35+
* indices placed together.
36+
* @param[in] m Number of selected points (length of idx).
37+
*
38+
* @return The number of points in the first subset within idx.
39+
*
40+
*/
41+
int _maxmin_divide(const double* grid, int* idx, int m) {
42+
assert(m > 1);
43+
if (m == 2) {
44+
return 1;
45+
}
46+
47+
std::vector<double> centroid(3, 0.0);
48+
for (int i = 0; i < m; ++i) {
49+
int j = idx[i];
50+
centroid[0] += grid[3*j ];
51+
centroid[1] += grid[3*j + 1];
52+
centroid[2] += grid[3*j + 2];
53+
}
54+
centroid[0] /= m;
55+
centroid[1] /= m;
56+
centroid[2] /= m;
57+
58+
// positions w.r.t. the centroid
59+
std::vector<double> R(3*m, 0.0);
60+
for (int i = 0; i < m; ++i) {
61+
int j = idx[i];
62+
R[3*i ] = grid[3*j ] - centroid[0];
63+
R[3*i + 1] = grid[3*j + 1] - centroid[1];
64+
R[3*i + 2] = grid[3*j + 2] - centroid[2];
65+
}
66+
67+
// The normal vector of the cut plane is taken to be the eigenvector
68+
// corresponding to the largest eigenvalue of the 3x3 matrix A = R*R^T.
69+
std::vector<double> A(9, 0.0);
70+
int i3 = 3, i1 = 1;
71+
double d0 = 0.0, d1 = 1.0;
72+
dsyrk_("U", "N", &i3, &m, &d1, R.data(), &i3, &d0, A.data(), &i3);
73+
74+
int info = 0, lwork = 102 /* determined by a work space query */;
75+
std::vector<double> e(3), work(lwork);
76+
dsyev_("V", "U", &i3, A.data(), &i3, e.data(), work.data(), &lwork, &info);
77+
double* n = A.data() + 6; // normal vector of the cut plane
78+
79+
// Rearrange the indices to put points in each subset together by
80+
// examining the signed distances of points to the cut plane (R^T*n).
81+
std::vector<double> dist(m);
82+
dgemv_("T", &i3, &m, &d1, R.data(), &i3, n, &i1, &d0, dist.data(), &i1);
83+
84+
int *head = idx;
85+
std::reverse_iterator<int*> tail(idx + m), rend(idx);
86+
auto is_negative = [&dist, &idx](int& j) { return dist[&j - idx] < 0; };
87+
while ( ( head = std::find_if(head, idx + m, is_negative) ) <
88+
( tail = std::find_if_not(tail, rend, is_negative) ).base() ) {
89+
std::swap(*head, *tail);
90+
std::swap(dist[head - idx], dist[tail.base() - idx - 1]);
91+
++head;
92+
++tail;
93+
}
94+
95+
return head - idx;
96+
}
97+
98+
} // end of anonymous namespace
99+
100+
101+
std::vector<int> Grid::Batch::maxmin(
102+
const double* grid,
103+
int* idx,
104+
int m,
105+
int m_thr
106+
) {
107+
if (m <= m_thr) {
108+
return std::vector<int>{0};
109+
}
110+
111+
int m_left = _maxmin_divide(grid, idx, m);
112+
113+
std::vector<int> left = maxmin(grid, idx, m_left, m_thr);
114+
std::vector<int> right = maxmin(grid, idx + m_left, m - m_left, m_thr);
115+
std::for_each(right.begin(), right.end(),
116+
[m_left](int& x) { x += m_left; }
117+
);
118+
119+
left.insert(left.end(), right.begin(), right.end());
120+
return left;
121+
}
122+
123+

‎source/module_base/grid/batch.h‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#ifndef GRID_BATCH_H
2+
#define GRID_BATCH_H
3+
4+
#include <vector>
5+
6+
namespace Grid {
7+
namespace Batch {
8+
9+
/**
10+
* @brief Divide a set of points into batches by the "MaxMin" algorithm.
11+
*
12+
* This function recursively uses cut planes to divide grid points into
13+
* two subsets using the "MaxMin" algorithm, until the number of points
14+
* in each subset (batch) is no more than m_thr.
15+
*
16+
* @param[in] grid Coordinates of all grid points.
17+
* grid[3*j], grid[3*j+1], grid[3*j+2] are
18+
* the x, y, z coordinates of the j-th point.
19+
* @param[in,out] idx Indices of the initial set within grid.
20+
* On return, idx will be rearranged such
21+
* that points belonging to the same batch
22+
* have their indices placed together.
23+
* @param[in] m Number of points in the initial set.
24+
* (length of idx)
25+
* @param[in] m_thr Size limit of a batch.
26+
*
27+
* @return Indices (for idx) of the first point in each batch.
28+
*
29+
* For example, given grid (illustrated by their indices) located as follows:
30+
*
31+
* 0 1 16 2 3 18
32+
* 4 5 17 6 7
33+
*
34+
*
35+
* 8 9 20 10 11
36+
* 12 13 19 14 15
37+
*
38+
* a possible outcome with m_thr = 4 and idx(in) = {0, 1, 2, ..., 15}
39+
* (idx may correspond to a subset of grid and does not have to be sorted,
40+
* but it must not contain duplicates) is:
41+
*
42+
* idx(out): 0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15
43+
* return : {0, 4, 8, 12}
44+
*
45+
* which means the selected set (labeled 0-15) is divided into 4 batches:
46+
* {0, 1, 4, 5}, {8, 9, 12, 13}, {2, 3, 6, 7}, {10, 11, 14, 15}.
47+
*
48+
*/
49+
std::vector<int> maxmin(const double* grid, int* idx, int m, int m_thr);
50+
51+
} // end of namespace Batch
52+
} // end of namespace Grid
53+
54+
#endif

‎source/module_base/grid/test/CMakeLists.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,3 +20,10 @@ AddTest(
2020
../radial.cpp
2121
../delley.cpp
2222
)
23+
24+
AddTest(
25+
TARGET test_batch
26+
SOURCES test_batch.cpp
27+
../batch.cpp
28+
LIBS ${math_libs}
29+
)

0 commit comments

Comments
 (0)