|
| 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 | + |
0 commit comments