From c33cb2caff952e7acfbea93ed6dfdf9d1cb6253d Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Fri, 22 Nov 2024 02:09:56 +0800
Subject: [PATCH 1/7] initial commit

---
 source/module_base/grid/batch.h | 74 +++++++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)
 create mode 100644 source/module_base/grid/batch.h

diff --git a/source/module_base/grid/batch.h b/source/module_base/grid/batch.h
new file mode 100644
index 0000000000..939b690f6e
--- /dev/null
+++ b/source/module_base/grid/batch.h
@@ -0,0 +1,74 @@
+#ifndef GRID_BATCH_H
+#define GRID_BATCH_H
+
+#include <vector>
+
+#include "module_base/blas_connector.h"
+#include "module_base/lapack_connector.h"
+#include <algorithm>
+#include <cassert>
+
+namespace Grid {
+namespace Batch {
+
+int _maxmin_core(int n, const double* grid, int* idx) {
+    assert(n > 1);
+    if (n == 2) {
+        return 1;
+    }
+
+    // center of the grid
+    std::vector<double> center(3, 0.0);
+    for (int i = 0; i < n; ++i) {
+        int ig = idx[i];
+        center[0] += grid[3*ig    ];
+        center[1] += grid[3*ig + 1];
+        center[2] += grid[3*ig + 2];
+    }
+    center[0] /= n;
+    center[1] /= n;
+    center[2] /= n;
+
+    // positions w.r.t. the center
+    std::vector<double> R(3*n, 0.0);
+    for (int i = 0; i < n; ++i) {
+        int ig = idx[i];
+        R[3*i    ] = grid[3*ig    ] - center[0];
+        R[3*i + 1] = grid[3*ig + 1] - center[1];
+        R[3*i + 2] = grid[3*ig + 2] - center[2];
+    }
+
+    // R*R^T
+    std::vector<double> A(9, 0.0);
+    int three_i = 3;
+    double one_d = 1.0, zero_d = 0.0;
+    dgemm_("N", "T", &three_i, &three_i, &n, &one_d, R.data(), &three_i, R.data(), &n, &zero_d, A.data(), &three_i);
+
+
+	//void dsyev_(const char* jobz,const char* uplo,const int* n,double *a,
+    //            const int* lda,double* w,double* work,const int* lwork, int* info);
+}
+
+
+std::vector<int> maxmin(int n_max, int n, const double* grid, int* idx) {
+    if (n <= n_max) {
+        return std::vector<int>{};
+    }
+
+    int n_left = _maxmin_core(n, grid, idx);
+
+    std::vector<int> delim_left = maxmin(n_max, n_left, grid, idx);
+    std::vector<int> delim_right = maxmin(n_max, n - n_left, grid + n_left, idx + n_left);;
+    std::for_each(delim_right.begin(), delim_right.end(), [n_left](int& x) { x += n_left; });
+
+    // merge all delimiters into delim_left
+    delim_left.resize(delim_left.size() + delim_right.size() + 1);
+    delim_left[delim_left.size()] = n_left;
+    std::copy(delim_right.begin(), delim_right.end(), delim_left.begin() + delim_left.size() + 1);
+    return delim_left;
+}
+
+} // end of namespace Batch
+} // end of namespace Grid
+
+#endif

From 1a2ba2d52096fd41630b7e69b8510437cedfe061 Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Fri, 22 Nov 2024 16:23:35 +0800
Subject: [PATCH 2/7] continue...

---
 source/module_base/grid/batch.cpp | 127 ++++++++++++++++++++++++++++++
 source/module_base/grid/batch.h   |  62 +--------------
 2 files changed, 128 insertions(+), 61 deletions(-)
 create mode 100644 source/module_base/grid/batch.cpp

diff --git a/source/module_base/grid/batch.cpp b/source/module_base/grid/batch.cpp
new file mode 100644
index 0000000000..73947302b0
--- /dev/null
+++ b/source/module_base/grid/batch.cpp
@@ -0,0 +1,127 @@
+#include "module_base/grid/batch.h"
+
+#include "module_base/blas_connector.h"
+#include "module_base/lapack_connector.h"
+#include <algorithm>
+#include <cassert>
+#include <iterator>
+
+namespace {
+
+/**
+ * @brief Bisect a set of points by the "MaxMin" algorithm.
+ *
+ * Given a selected set of grid points specified by the indices `idx`,
+ * bisect this set by a cut plane {x|n^T*(x-c) = 0} where the normal
+ * vector n and the point c are determined by the "MaxMin" problem:
+ *
+ *      max min sum_{i=1}^{m} [n^T*(r[idx[i]] - c)]^2
+ *       n   c
+ *
+ * here r[j] = (grid[3*j], grid[3*j+1], grid[3*j+2]) is the position of
+ * the j-th point.
+ *
+ * It can be shown that the optimal c is the centroid of the points, and
+ * the optimal n is the eigenvector corresponding to the largest eigenvalue
+ * of the matrix R*R^T, where R is the matrix whose i-th column is
+ * r[idx[i]] - c.
+ *
+ * param[in]        m       Number of the selected points (size of idx).
+ * param[in]        grid    Coordinates of all grid points.
+ *                          grid[3*j], grid[3*j+1], grid[3*j+2] are the
+ *                          x, y, z coordinates of the j-th point.
+ *                          The size of grid is at least 3*m.
+ * param[in,out]    idx     Indices of the selected points within grid.
+ *
+ */
+int _maxmin_bisect(int m, const double* grid, int* idx) {
+    assert(m > 1);
+    if (m == 2) {
+        return 1;
+    }
+
+    std::vector<double> centroid(3, 0.0);
+    for (int i = 0; i < m; ++i) {
+        int ig = idx[i];
+        centroid[0] += grid[3*ig    ];
+        centroid[1] += grid[3*ig + 1];
+        centroid[2] += grid[3*ig + 2];
+    }
+    centroid[0] /= m;
+    centroid[1] /= m;
+    centroid[2] /= m;
+
+    // positions w.r.t. the centroid
+    std::vector<double> R(3*m, 0.0);
+    for (int i = 0; i < m; ++i) {
+        int j = idx[i];
+        R[3*i    ] = grid[3*j    ] - centroid[0];
+        R[3*i + 1] = grid[3*j + 1] - centroid[1];
+        R[3*i + 2] = grid[3*j + 2] - centroid[2];
+    }
+
+    // A = R*R^T is a 3-by-3 matrix
+    std::vector<double> A(9, 0.0);
+    int i3 = 3, i1 = 1;
+    double d0 = 0.0, d1 = 1.0;
+    dsyrk_("U", "N", &i3, &m, &d1, R.data(), &i3, &d0, A.data(), &i3);
+
+    // eigenpairs of A
+    int info = 0, lwork = 102 /* determined by a work space query */;
+    std::vector<double> e(3), work(lwork);
+    dsyev_("V", "U", &i3, A.data(), &i3, e.data(), work.data(), &lwork, &info);
+
+    // normal vector of the cut plane
+    // (eigenvector corresponding to the largest eigenvalue)
+    double* n = A.data() + 6;
+
+    // (signed) distance w.r.t. the cut plane
+    std::vector<double> dist(m);
+    for (int i = 0; i < m; ++i) {
+        dist[i] = ddot_(&i3, R.data() + 3*i, &i1, n, &i1);
+    }
+
+    // 
+    int *head = idx;
+    std::reverse_iterator<int*> tail(idx + m), rend(idx);
+    auto is_negative = [&dist](int j) { return dist[j] < 0; };
+    auto is_nonnegative = [&dist](int j) { return dist[j] >= 0; };
+    while ( (head = std::find(head, idx + m, is_negative)) <
+            (tail = std::find(tail, rend, is_nonnegative)).base()) {
+        std::swap(*head, *tail);
+        std::swap(dist[head - idx], dist[tail.base() - idx]);
+        ++head;
+        ++tail;
+    }
+
+    return std::find(idx, idx + m, is_nonnegative) - idx;
+}
+
+} // end of anonymous namespace
+
+std::vector<int> Grid::Batch::maxmin(
+    int n_max,
+    int n,
+    const double* grid,
+    int* idx
+) {
+    if (n <= n_max) {
+        return std::vector<int>{};
+    }
+
+    int n_left = _maxmin_bisect(n, grid, idx);
+
+    std::vector<int> delim_left = maxmin(n_max, n_left, grid, idx);
+    std::vector<int> delim_right = maxmin(n_max, n - n_left, grid + n_left, idx + n_left);
+    std::for_each(delim_right.begin(), delim_right.end(),
+        [n_left](int& x) { x += n_left; }
+    );
+
+    // merge all delimiters into delim_left
+    delim_left.resize(delim_left.size() + delim_right.size() + 1);
+    delim_left[delim_left.size()] = n_left;
+    std::copy(delim_right.begin(), delim_right.end(), delim_left.begin() + delim_left.size() + 1);
+    return delim_left;
+}
+
+
diff --git a/source/module_base/grid/batch.h b/source/module_base/grid/batch.h
index 939b690f6e..c72936e65e 100644
--- a/source/module_base/grid/batch.h
+++ b/source/module_base/grid/batch.h
@@ -3,70 +3,10 @@
 
 #include <vector>
 
-#include "module_base/blas_connector.h"
-#include "module_base/lapack_connector.h"
-#include <algorithm>
-#include <cassert>
-
 namespace Grid {
 namespace Batch {
 
-int _maxmin_core(int n, const double* grid, int* idx) {
-    assert(n > 1);
-    if (n == 2) {
-        return 1;
-    }
-
-    // center of the grid
-    std::vector<double> center(3, 0.0);
-    for (int i = 0; i < n; ++i) {
-        int ig = idx[i];
-        center[0] += grid[3*ig    ];
-        center[1] += grid[3*ig + 1];
-        center[2] += grid[3*ig + 2];
-    }
-    center[0] /= n;
-    center[1] /= n;
-    center[2] /= n;
-
-    // positions w.r.t. the center
-    std::vector<double> R(3*n, 0.0);
-    for (int i = 0; i < n; ++i) {
-        int ig = idx[i];
-        R[3*i    ] = grid[3*ig    ] - center[0];
-        R[3*i + 1] = grid[3*ig + 1] - center[1];
-        R[3*i + 2] = grid[3*ig + 2] - center[2];
-    }
-
-    // R*R^T
-    std::vector<double> A(9, 0.0);
-    int three_i = 3;
-    double one_d = 1.0, zero_d = 0.0;
-    dgemm_("N", "T", &three_i, &three_i, &n, &one_d, R.data(), &three_i, R.data(), &n, &zero_d, A.data(), &three_i);
-
-
-	//void dsyev_(const char* jobz,const char* uplo,const int* n,double *a,
-    //            const int* lda,double* w,double* work,const int* lwork, int* info);
-}
-
-
-std::vector<int> maxmin(int n_max, int n, const double* grid, int* idx) {
-    if (n <= n_max) {
-        return std::vector<int>{};
-    }
-
-    int n_left = _maxmin_core(n, grid, idx);
-
-    std::vector<int> delim_left = maxmin(n_max, n_left, grid, idx);
-    std::vector<int> delim_right = maxmin(n_max, n - n_left, grid + n_left, idx + n_left);;
-    std::for_each(delim_right.begin(), delim_right.end(), [n_left](int& x) { x += n_left; });
-
-    // merge all delimiters into delim_left
-    delim_left.resize(delim_left.size() + delim_right.size() + 1);
-    delim_left[delim_left.size()] = n_left;
-    std::copy(delim_right.begin(), delim_right.end(), delim_left.begin() + delim_left.size() + 1);
-    return delim_left;
-}
+std::vector<int> maxmin(int n_max, int n, const double* grid, int* idx);
 
 } // end of namespace Batch
 } // end of namespace Grid

From d3ee81306cd472ed6a5d522cc81c9764e293d988 Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Fri, 22 Nov 2024 19:38:54 +0800
Subject: [PATCH 3/7] crude test...

---
 source/module_base/grid/batch.cpp           |  89 +++++++-------
 source/module_base/grid/batch.h             |   8 ++
 source/module_base/grid/test/test_batch.cpp | 127 ++++++++++++++++++++
 3 files changed, 181 insertions(+), 43 deletions(-)
 create mode 100644 source/module_base/grid/test/test_batch.cpp

diff --git a/source/module_base/grid/batch.cpp b/source/module_base/grid/batch.cpp
index 73947302b0..1876ec0c83 100644
--- a/source/module_base/grid/batch.cpp
+++ b/source/module_base/grid/batch.cpp
@@ -1,21 +1,22 @@
 #include "module_base/grid/batch.h"
 
-#include "module_base/blas_connector.h"
-#include "module_base/lapack_connector.h"
 #include <algorithm>
 #include <cassert>
 #include <iterator>
 
+#include "module_base/blas_connector.h"
+#include "module_base/lapack_connector.h"
+
 namespace {
 
 /**
- * @brief Bisect a set of points by the "MaxMin" algorithm.
+ * @brief Divide a set of points into two subsets by the "MaxMin" algorithm.
  *
- * Given a selected set of grid points specified by the indices `idx`,
- * bisect this set by a cut plane {x|n^T*(x-c) = 0} where the normal
- * vector n and the point c are determined by the "MaxMin" problem:
+ * This function divides a given set of grid points by a cut plane
+ * {x|n^T*(x-c) = 0} where the normal vector n and the point c are
+ * determined by the "MaxMin" problem:
  *
- *      max min sum_{i=1}^{m} [n^T*(r[idx[i]] - c)]^2
+ *      max min sum_{i=1}^{m} [n^T * (r[idx[i]] - c)]^2
  *       n   c
  *
  * here r[j] = (grid[3*j], grid[3*j+1], grid[3*j+2]) is the position of
@@ -26,15 +27,19 @@ namespace {
  * of the matrix R*R^T, where R is the matrix whose i-th column is
  * r[idx[i]] - c.
  *
- * param[in]        m       Number of the selected points (size of idx).
- * param[in]        grid    Coordinates of all grid points.
+ * @param[in]       m       Number of selected points (length of idx).
+ * @param[in]       grid    Coordinates of all grid points.
  *                          grid[3*j], grid[3*j+1], grid[3*j+2] are the
  *                          x, y, z coordinates of the j-th point.
- *                          The size of grid is at least 3*m.
- * param[in,out]    idx     Indices of the selected points within grid.
+ *                          The length of grid is at least 3*m.
+ * @param[in,out]   idx     Indices of the selected points within grid.
+ *                          On exit, the indices are rearranged such that
+ *                          points in each subset are put together.
+ *
+ * @return The number of points in the first subset (according to idx).
  *
  */
-int _maxmin_bisect(int m, const double* grid, int* idx) {
+int _maxmin_divide(int m, const double* grid, int* idx) {
     assert(m > 1);
     if (m == 2) {
         return 1;
@@ -42,10 +47,10 @@ int _maxmin_bisect(int m, const double* grid, int* idx) {
 
     std::vector<double> centroid(3, 0.0);
     for (int i = 0; i < m; ++i) {
-        int ig = idx[i];
-        centroid[0] += grid[3*ig    ];
-        centroid[1] += grid[3*ig + 1];
-        centroid[2] += grid[3*ig + 2];
+        int j = idx[i];
+        centroid[0] += grid[3*j    ];
+        centroid[1] += grid[3*j + 1];
+        centroid[2] += grid[3*j + 2];
     }
     centroid[0] /= m;
     centroid[1] /= m;
@@ -60,68 +65,66 @@ int _maxmin_bisect(int m, const double* grid, int* idx) {
         R[3*i + 2] = grid[3*j + 2] - centroid[2];
     }
 
-    // A = R*R^T is a 3-by-3 matrix
+    // The normal vector of the cut plane is taken to be the eigenvector
+    // corresponding to the largest eigenvalue of the 3x3 matrix A = R*R^T.
     std::vector<double> A(9, 0.0);
     int i3 = 3, i1 = 1;
     double d0 = 0.0, d1 = 1.0;
     dsyrk_("U", "N", &i3, &m, &d1, R.data(), &i3, &d0, A.data(), &i3);
 
-    // eigenpairs of A
     int info = 0, lwork = 102 /* determined by a work space query */;
     std::vector<double> e(3), work(lwork);
     dsyev_("V", "U", &i3, A.data(), &i3, e.data(), work.data(), &lwork, &info);
+    double* n = A.data() + 6; // normal vector of the cut plane
 
-    // normal vector of the cut plane
-    // (eigenvector corresponding to the largest eigenvalue)
-    double* n = A.data() + 6;
-
-    // (signed) distance w.r.t. the cut plane
+    // Rearrange the indices to put points in each subset together by
+    // examining the signed distances of the points to the cut plane (R^T*n).
     std::vector<double> dist(m);
-    for (int i = 0; i < m; ++i) {
-        dist[i] = ddot_(&i3, R.data() + 3*i, &i1, n, &i1);
-    }
+    dgemv_("T", &i3, &m, &d1, R.data(), &i3, n, &i1, &d0, dist.data(), &i1);
 
-    // 
     int *head = idx;
     std::reverse_iterator<int*> tail(idx + m), rend(idx);
     auto is_negative = [&dist](int j) { return dist[j] < 0; };
     auto is_nonnegative = [&dist](int j) { return dist[j] >= 0; };
-    while ( (head = std::find(head, idx + m, is_negative)) <
-            (tail = std::find(tail, rend, is_nonnegative)).base()) {
+    while ( ( head = std::find(head, idx + m, is_negative) ) <
+            ( tail = std::find(tail, rend, is_nonnegative) ).base() ) {
         std::swap(*head, *tail);
         std::swap(dist[head - idx], dist[tail.base() - idx]);
         ++head;
         ++tail;
     }
 
-    return std::find(idx, idx + m, is_nonnegative) - idx;
+    return std::find(idx, idx + m, is_negative) - idx;
 }
 
 } // end of anonymous namespace
 
+
 std::vector<int> Grid::Batch::maxmin(
-    int n_max,
-    int n,
+    int m_max,
+    int m,
     const double* grid,
     int* idx
 ) {
-    if (n <= n_max) {
+    if (m <= m_max) {
         return std::vector<int>{};
     }
 
-    int n_left = _maxmin_bisect(n, grid, idx);
+    int m_left = _maxmin_divide(m, grid, idx);
 
-    std::vector<int> delim_left = maxmin(n_max, n_left, grid, idx);
-    std::vector<int> delim_right = maxmin(n_max, n - n_left, grid + n_left, idx + n_left);
-    std::for_each(delim_right.begin(), delim_right.end(),
-        [n_left](int& x) { x += n_left; }
+    // recursively divide the subsets
+    std::vector<int> left = maxmin(m_max, m_left, grid, idx);
+    std::vector<int> right = maxmin(m_max, m - m_left, grid, idx + m_left);
+    std::for_each(right.begin(), right.end(),
+        [m_left](int& x) { x += m_left; }
     );
 
-    // merge all delimiters into delim_left
-    delim_left.resize(delim_left.size() + delim_right.size() + 1);
-    delim_left[delim_left.size()] = n_left;
-    std::copy(delim_right.begin(), delim_right.end(), delim_left.begin() + delim_left.size() + 1);
-    return delim_left;
+    // merge all delimiters
+    int sz_left = left.size();
+    left.resize(sz_left + right.size() + 1);
+    left[sz_left] = m_left;
+    std::copy(right.begin(), right.end(), left.begin() + sz_left + 1);
+    return left;
 }
 
 
diff --git a/source/module_base/grid/batch.h b/source/module_base/grid/batch.h
index c72936e65e..04c9c2e57a 100644
--- a/source/module_base/grid/batch.h
+++ b/source/module_base/grid/batch.h
@@ -6,6 +6,14 @@
 namespace Grid {
 namespace Batch {
 
+/**
+ * @brief Divide a set of points into batches by the "MaxMin" algorithm.
+ *
+ * This function recursively divides a given set of grid points into two
+ * subsets by a cut plane using the "MaxMin" algorithm, until the number
+ * of points in each subset is no more than n_max.
+ *
+ */
 std::vector<int> maxmin(int n_max, int n, const double* grid, int* idx);
 
 } // end of namespace Batch
diff --git a/source/module_base/grid/test/test_batch.cpp b/source/module_base/grid/test/test_batch.cpp
new file mode 100644
index 0000000000..f7d52fbc89
--- /dev/null
+++ b/source/module_base/grid/test/test_batch.cpp
@@ -0,0 +1,127 @@
+#include "module_base/grid/batch.h"
+
+#include "gtest/gtest.h"
+#include <algorithm>
+#include <random>
+
+#ifdef __MPI
+#include <mpi.h>
+#endif
+
+using namespace Grid::Batch;
+
+
+class BatchTest: public ::testing::Test
+{
+protected:
+    void SetUp();
+
+    std::vector<double> grid_;
+    std::vector<int> idx_;
+
+    int n_each_ = 10;
+    double offset_ = 10.0;
+    double width_ = 1.0;
+};
+
+std::vector<double> gen_octant_cluster(int n_each, double offset, double width) {
+
+    // Generates a set of points consisting of 8 well-separated, equal-sized
+    // clusters located in individual octants.
+
+    std::vector<double> grid(n_each * 8);
+    int I = 0;
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_real_distribution<double> dis(-width, width);
+
+    for (int sign_x : {-1, 1}) {
+        for (int sign_y : {-1, 1}) {
+            for (int sign_z : {-1, 1}) {
+                for (int i = 0; i < n_each; ++i) {
+                    grid[3*I    ] = sign_x * offset + dis(gen);
+                    grid[3*I + 1] = sign_y * offset + dis(gen);
+                    grid[3*I + 2] = sign_z * offset + dis(gen);
+                    ++I;
+                }
+            }
+        }
+    }
+
+    return grid;
+}
+
+bool is_same_octant(int ngrid, const double* grid) {
+    if (ngrid == 0) {
+        return true;
+    }
+    bool is_positive_x = grid[0] > 0;
+    bool is_positive_y = grid[1] > 0;
+    bool is_positive_z = grid[2] > 0;
+    const double* end = grid + 3 * ngrid;
+    for (; grid != end; grid += 3) {
+        if ( is_positive_x != (grid[0] > 0) ||
+             is_positive_y != (grid[1] > 0) ||
+             is_positive_z != (grid[2] > 0) ) {
+            return false;
+        }
+    }
+    return true;
+}
+
+
+void BatchTest::SetUp()
+{
+    grid_ = gen_octant_cluster(n_each_, offset_, width_);
+
+    idx_.resize(grid_.size());
+    std::iota(idx_.begin(), idx_.end(), 0);
+
+    std::random_device rd;
+    std::mt19937 g(rd());
+    std::shuffle(idx_.begin(), idx_.end(), g);
+}
+
+
+TEST_F(BatchTest, MaxMinOctantCluster)
+{
+    // This test applies maxmin to a set of points consisting of 8
+    // well-separated, equal-sized clusters located in individual octants.
+    // The resulting batches should be able to recover this structure.
+
+    std::vector<int> delim = 
+        maxmin(n_each_, grid_.size(), grid_.data(), idx_.data());
+
+    EXPECT_EQ(delim.size(), 7);
+    for (int i = 0; i < 7; ++i) {
+        // check number of points in each batch via index delimiters
+        EXPECT_EQ(delim[i], (i+1) * n_each_);
+
+        // verify that points in each batch is in the same octant
+        std::vector<double> batch(3 * n_each_);
+        for (int j = 0; j < n_each_; ++j) {
+            for (int k = 0; k < 3; ++k) {
+                batch[3*j + k] = grid_[3*(i*n_each_ + j) + k];
+            }
+        }
+        EXPECT_TRUE(is_same_octant(n_each_, batch.data()));
+    }
+}
+
+
+int main(int argc, char** argv)
+{
+#ifdef __MPI
+    MPI_Init(&argc, &argv);
+#endif
+
+    testing::InitGoogleTest(&argc, argv);
+    int result = RUN_ALL_TESTS();
+
+#ifdef __MPI
+    MPI_Finalize();
+#endif
+
+    return result;
+}

From 0369832696618cc1ba0c31c1a4cbd20a63c95d71 Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Sat, 23 Nov 2024 18:08:58 +0800
Subject: [PATCH 4/7] continue...

---
 source/module_base/blas_connector.h | 16 +++++++++++++++-
 source/module_base/grid/batch.cpp   | 28 +++++++++++-----------------
 source/module_base/grid/batch.h     | 18 ++++++++++++++----
 3 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/source/module_base/blas_connector.h b/source/module_base/blas_connector.h
index e7b78bab06..090d8512ee 100644
--- a/source/module_base/blas_connector.h
+++ b/source/module_base/blas_connector.h
@@ -39,6 +39,20 @@ extern "C"
 	double dnrm2_( const int *n, const double *X, const int *incX );
 	double dznrm2_( const int *n, const std::complex<double> *X, const int *incX );
 
+    // symmetric rank-k update
+    void dsyrk_(
+        const char* uplo,
+        const char* trans,
+        const int* n,
+        const int* k,
+        const double* alpha,
+        const double* a,
+        const int* lda,
+        const double* beta,
+        double* c,
+        const int* ldc
+    );
+
 	// level 2: matrix-std::vector operations, O(n^2) data and O(n^2) work.
 	void sgemv_(const char*const transa, const int*const m, const int*const n,
 		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,
 */
 
 #endif // GATHER_INFO
-#endif // BLAS_CONNECTOR_H
\ No newline at end of file
+#endif // BLAS_CONNECTOR_H
diff --git a/source/module_base/grid/batch.cpp b/source/module_base/grid/batch.cpp
index 1876ec0c83..34ca268401 100644
--- a/source/module_base/grid/batch.cpp
+++ b/source/module_base/grid/batch.cpp
@@ -78,16 +78,15 @@ int _maxmin_divide(int m, const double* grid, int* idx) {
     double* n = A.data() + 6; // normal vector of the cut plane
 
     // Rearrange the indices to put points in each subset together by
-    // examining the signed distances of the points to the cut plane (R^T*n).
+    // examining the signed distances of points to the cut plane (R^T*n).
     std::vector<double> dist(m);
     dgemv_("T", &i3, &m, &d1, R.data(), &i3, n, &i1, &d0, dist.data(), &i1);
 
     int *head = idx;
     std::reverse_iterator<int*> tail(idx + m), rend(idx);
     auto is_negative = [&dist](int j) { return dist[j] < 0; };
-    auto is_nonnegative = [&dist](int j) { return dist[j] >= 0; };
-    while ( ( head = std::find(head, idx + m, is_negative) ) <
-            ( tail = std::find(tail, rend, is_nonnegative) ).base() ) {
+    while ( ( head = std::find_if(head, idx + m, is_negative) ) <
+            ( tail = std::find_if_not(tail, rend, is_negative) ).base() ) {
         std::swap(*head, *tail);
         std::swap(dist[head - idx], dist[tail.base() - idx]);
         ++head;
@@ -101,29 +100,24 @@ int _maxmin_divide(int m, const double* grid, int* idx) {
 
 
 std::vector<int> Grid::Batch::maxmin(
-    int m_max,
-    int m,
     const double* grid,
-    int* idx
+    int* idx,
+    int m,
+    int m_thr
 ) {
-    if (m <= m_max) {
-        return std::vector<int>{};
+    if (m <= m_thr) {
+        return std::vector<int>{0};
     }
 
     int m_left = _maxmin_divide(m, grid, idx);
 
-    // recursively divide the subsets
-    std::vector<int> left = maxmin(m_max, m_left, grid, idx);
-    std::vector<int> right = maxmin(m_max, m - m_left, grid, idx + m_left);
+    std::vector<int> left = maxmin(grid, idx, m_left, m_thr);
+    std::vector<int> right = maxmin(grid, idx + m_left, m - m_left, m_thr);
     std::for_each(right.begin(), right.end(),
         [m_left](int& x) { x += m_left; }
     );
 
-    // merge all delimiters
-    int sz_left = left.size();
-    left.resize(sz_left + right.size() + 1);
-    left[sz_left] = m_left;
-    std::copy(right.begin(), right.end(), left.begin() + sz_left + 1);
+    left.insert(left.end(), right.begin(), right.end());
     return left;
 }
 
diff --git a/source/module_base/grid/batch.h b/source/module_base/grid/batch.h
index 04c9c2e57a..7922667d65 100644
--- a/source/module_base/grid/batch.h
+++ b/source/module_base/grid/batch.h
@@ -9,12 +9,22 @@ namespace Batch {
 /**
  * @brief Divide a set of points into batches by the "MaxMin" algorithm.
  *
- * This function recursively divides a given set of grid points into two
- * subsets by a cut plane using the "MaxMin" algorithm, until the number
- * of points in each subset is no more than n_max.
+ * This function recursively uses a cut plane to divide a set of grid
+ * points into two subsets using the "MaxMin" algorithm, until the
+ * number of points in each subset (batch) is no more than m_thr.
+ *
+ * @param[in]       grid        Coordinates of all grid points.
+ * @param[in,out]   idx         Indices of the initial set within grid.
+ *                              On return, idx will be rearranged such
+ *                              that points belonging to the same subset
+ *                              are grouped together.
+ * @param[in]       m           Number of points in the initial set.
+ * @param[in]       m_thr       Size limit of subset.
+ *
+ * @return          Indices (within idx) of the first point in each batch.
  *
  */
-std::vector<int> maxmin(int n_max, int n, const double* grid, int* idx);
+std::vector<int> maxmin(const double* grid, int* idx, int m, int m_thr);
 
 } // end of namespace Batch
 } // end of namespace Grid

From f8b00a2ac43c3a1e1e13f1971b65ee17bb01ba83 Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Sat, 23 Nov 2024 23:02:18 +0800
Subject: [PATCH 5/7] complete!

---
 source/module_base/grid/batch.cpp           | 21 ++++----
 source/module_base/grid/batch.h             | 34 ++++++++++---
 source/module_base/grid/test/CMakeLists.txt |  7 +++
 source/module_base/grid/test/test_batch.cpp | 55 +++++++++++++--------
 4 files changed, 78 insertions(+), 39 deletions(-)

diff --git a/source/module_base/grid/batch.cpp b/source/module_base/grid/batch.cpp
index 34ca268401..26c9f4d60e 100644
--- a/source/module_base/grid/batch.cpp
+++ b/source/module_base/grid/batch.cpp
@@ -24,22 +24,21 @@ namespace {
  *
  * It can be shown that the optimal c is the centroid of the points, and
  * the optimal n is the eigenvector corresponding to the largest eigenvalue
- * of the matrix R*R^T, where R is the matrix whose i-th column is
- * r[idx[i]] - c.
+ * of the matrix R*R^T, where the i-th column of R is r[idx[i]] - c.
  *
- * @param[in]       m       Number of selected points (length of idx).
  * @param[in]       grid    Coordinates of all grid points.
  *                          grid[3*j], grid[3*j+1], grid[3*j+2] are the
  *                          x, y, z coordinates of the j-th point.
- *                          The length of grid is at least 3*m.
  * @param[in,out]   idx     Indices of the selected points within grid.
- *                          On exit, the indices are rearranged such that
- *                          points in each subset are put together.
+ *                          On return, idx will be rearranged such that
+ *                          points belonging to the same subset have their
+ *                          indices placed together.
+ * @param[in]       m       Number of selected points (length of idx).
  *
  * @return The number of points in the first subset (according to idx).
  *
  */
-int _maxmin_divide(int m, const double* grid, int* idx) {
+int _maxmin_divide(const double* grid, int* idx, int m) {
     assert(m > 1);
     if (m == 2) {
         return 1;
@@ -84,16 +83,16 @@ int _maxmin_divide(int m, const double* grid, int* idx) {
 
     int *head = idx;
     std::reverse_iterator<int*> tail(idx + m), rend(idx);
-    auto is_negative = [&dist](int j) { return dist[j] < 0; };
+    auto is_negative = [&dist, &idx](int& j) { return dist[&j - idx] < 0; };
     while ( ( head = std::find_if(head, idx + m, is_negative) ) <
             ( tail = std::find_if_not(tail, rend, is_negative) ).base() ) {
         std::swap(*head, *tail);
-        std::swap(dist[head - idx], dist[tail.base() - idx]);
+        std::swap(dist[head - idx], dist[tail.base() - idx - 1]);
         ++head;
         ++tail;
     }
 
-    return std::find(idx, idx + m, is_negative) - idx;
+    return head - idx;
 }
 
 } // end of anonymous namespace
@@ -109,7 +108,7 @@ std::vector<int> Grid::Batch::maxmin(
         return std::vector<int>{0};
     }
 
-    int m_left = _maxmin_divide(m, grid, idx);
+    int m_left = _maxmin_divide(grid, idx, m);
 
     std::vector<int> left = maxmin(grid, idx, m_left, m_thr);
     std::vector<int> right = maxmin(grid, idx + m_left, m - m_left, m_thr);
diff --git a/source/module_base/grid/batch.h b/source/module_base/grid/batch.h
index 7922667d65..a2646b9b70 100644
--- a/source/module_base/grid/batch.h
+++ b/source/module_base/grid/batch.h
@@ -9,19 +9,39 @@ namespace Batch {
 /**
  * @brief Divide a set of points into batches by the "MaxMin" algorithm.
  *
- * This function recursively uses a cut plane to divide a set of grid
- * points into two subsets using the "MaxMin" algorithm, until the
- * number of points in each subset (batch) is no more than m_thr.
+ * This function recursively uses cut planes to divide grid points into
+ * two subsets using the "MaxMin" algorithm, until the number of points
+ * in each subset (batch) is no more than m_thr.
  *
  * @param[in]       grid        Coordinates of all grid points.
+ *                              grid[3*j], grid[3*j+1], grid[3*j+2] are
+ *                              the x, y, z coordinates of the j-th point.
  * @param[in,out]   idx         Indices of the initial set within grid.
  *                              On return, idx will be rearranged such
- *                              that points belonging to the same subset
- *                              are grouped together.
+ *                              that points belonging to the same batch
+ *                              have their indices placed together.
  * @param[in]       m           Number of points in the initial set.
- * @param[in]       m_thr       Size limit of subset.
+ *                              (length of idx)
+ * @param[in]       m_thr       Size limit of a batch.
  *
- * @return          Indices (within idx) of the first point in each batch.
+ * @return          Indices (for idx) of the first point in each batch.
+ *
+ * For example, given grid (illustrated by their indices) located as follows:
+ *
+ *      0  1  16          2  3  18
+ *      4  5  17            6  7
+ *
+ *
+ *      8  9             20 10 11
+ *     12 13 19           14 15
+ *
+ * a possible outcome with m_thr = 4 and idx(in) = {0, 1, 2, ..., 15} is:
+ *
+ * idx(out): 0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15
+ * return  : {0, 4, 8, 12}
+ *
+ * which means the selected set (labeled 0-15) is divided into 4 batches:
+ * {0, 1, 4, 5}, {8, 9, 12, 13}, {2, 3, 6, 7}, {10, 11, 14, 15}.
  *
  */
 std::vector<int> maxmin(const double* grid, int* idx, int m, int m_thr);
diff --git a/source/module_base/grid/test/CMakeLists.txt b/source/module_base/grid/test/CMakeLists.txt
index e403c8c2c6..a0f1b6fa9f 100644
--- a/source/module_base/grid/test/CMakeLists.txt
+++ b/source/module_base/grid/test/CMakeLists.txt
@@ -20,3 +20,10 @@ AddTest(
   ../radial.cpp
   ../delley.cpp
 )
+
+AddTest(
+  TARGET test_batch
+  SOURCES test_batch.cpp
+  ../batch.cpp
+  LIBS ${math_libs}
+)
diff --git a/source/module_base/grid/test/test_batch.cpp b/source/module_base/grid/test/test_batch.cpp
index f7d52fbc89..6ba1ae2161 100644
--- a/source/module_base/grid/test/test_batch.cpp
+++ b/source/module_base/grid/test/test_batch.cpp
@@ -10,7 +10,6 @@
 
 using namespace Grid::Batch;
 
-
 class BatchTest: public ::testing::Test
 {
 protected:
@@ -19,17 +18,28 @@ class BatchTest: public ::testing::Test
     std::vector<double> grid_;
     std::vector<int> idx_;
 
+    // parameters for cluster generation
     int n_each_ = 10;
-    double offset_ = 10.0;
     double width_ = 1.0;
+
+    // These offsets should be different from each other as maxmin might
+    // fail for highly symmetric, well-separated clusters.
+    // Consider the case where the 8 clusters as a whole have octahedral
+    // symmetry. In this case, R*R^T must be proprotional to the identity,
+    // and eigenvalues are three-fold degenerate, because xy, yz and zx
+    // plane are equivalent in terms of the maxmin optimization problem.
+    // This means eigenvectors are arbitrary in this case.
+    double offset_x_ = 7.0;
+    double offset_y_ = 8.0;
+    double offset_z_ = 9.0;
 };
 
-std::vector<double> gen_octant_cluster(int n_each, double offset, double width) {
+std::vector<double> gen_octant_cluster(int n_each, double offset_x, double offset_y, double offset_z, double width) {
 
     // Generates a set of points consisting of 8 well-separated, equal-sized
     // clusters located in individual octants.
 
-    std::vector<double> grid(n_each * 8);
+    std::vector<double> grid(n_each * 8 * 3);
     int I = 0;
 
     std::random_device rd;
@@ -40,15 +50,14 @@ std::vector<double> gen_octant_cluster(int n_each, double offset, double width)
         for (int sign_y : {-1, 1}) {
             for (int sign_z : {-1, 1}) {
                 for (int i = 0; i < n_each; ++i) {
-                    grid[3*I    ] = sign_x * offset + dis(gen);
-                    grid[3*I + 1] = sign_y * offset + dis(gen);
-                    grid[3*I + 2] = sign_z * offset + dis(gen);
+                    grid[3*I    ] = sign_x * offset_x + dis(gen);
+                    grid[3*I + 1] = sign_y * offset_y + dis(gen);
+                    grid[3*I + 2] = sign_z * offset_z + dis(gen);
                     ++I;
                 }
             }
         }
     }
-
     return grid;
 }
 
@@ -73,9 +82,9 @@ bool is_same_octant(int ngrid, const double* grid) {
 
 void BatchTest::SetUp()
 {
-    grid_ = gen_octant_cluster(n_each_, offset_, width_);
+    grid_ = gen_octant_cluster(n_each_, offset_x_, offset_y_, offset_z_, width_);
 
-    idx_.resize(grid_.size());
+    idx_.resize(grid_.size() / 3);
     std::iota(idx_.begin(), idx_.end(), 0);
 
     std::random_device rd;
@@ -91,21 +100,25 @@ TEST_F(BatchTest, MaxMinOctantCluster)
     // The resulting batches should be able to recover this structure.
 
     std::vector<int> delim = 
-        maxmin(n_each_, grid_.size(), grid_.data(), idx_.data());
+        maxmin(grid_.data(), idx_.data(), grid_.size() / 3, n_each_);
+
+    EXPECT_EQ(delim.size(), 8);
 
-    EXPECT_EQ(delim.size(), 7);
-    for (int i = 0; i < 7; ++i) {
-        // check number of points in each batch via index delimiters
-        EXPECT_EQ(delim[i], (i+1) * n_each_);
+    std::vector<double> grid_batch(3 * n_each_);
+    for (int i = 0; i < 8; ++i) {
 
-        // verify that points in each batch is in the same octant
-        std::vector<double> batch(3 * n_each_);
+        EXPECT_EQ(delim[i], i * n_each_);
+
+        // collect points within the present batch
         for (int j = 0; j < n_each_; ++j) {
-            for (int k = 0; k < 3; ++k) {
-                batch[3*j + k] = grid_[3*(i*n_each_ + j) + k];
-            }
+            int ig = idx_[delim[i] + j];
+            grid_batch[3*j    ] = grid_[3*ig    ];
+            grid_batch[3*j + 1] = grid_[3*ig + 1];
+            grid_batch[3*j + 2] = grid_[3*ig + 2];
         }
-        EXPECT_TRUE(is_same_octant(n_each_, batch.data()));
+
+        // verify that points in a batch reside in the same octant
+        EXPECT_TRUE(is_same_octant(n_each_, grid_batch.data()));
     }
 }
 

From 4c8876e56c8a36acf6f29a3f286a97f8b4f3265d Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Mon, 25 Nov 2024 13:59:43 +0800
Subject: [PATCH 6/7] more test

---
 source/module_base/grid/test/test_batch.cpp | 125 +++++++++++++++-----
 1 file changed, 93 insertions(+), 32 deletions(-)

diff --git a/source/module_base/grid/test/test_batch.cpp b/source/module_base/grid/test/test_batch.cpp
index 6ba1ae2161..d78446db4d 100644
--- a/source/module_base/grid/test/test_batch.cpp
+++ b/source/module_base/grid/test/test_batch.cpp
@@ -10,64 +10,112 @@
 
 using namespace Grid::Batch;
 
+
 class BatchTest: public ::testing::Test
 {
 protected:
-    void SetUp();
 
     std::vector<double> grid_;
     std::vector<int> idx_;
 
-    // parameters for cluster generation
-    int n_each_ = 10;
-    double width_ = 1.0;
-
-    // These offsets should be different from each other as maxmin might
-    // fail for highly symmetric, well-separated clusters.
+    // parameters for octant clusters
+    const int n_batch_oct_ = 10;
+    const double width_oct_ = 1.0;
+    const double offset_x_ = 7.0;
+    const double offset_y_ = 8.0;
+    const double offset_z_ = 9.0;
+    // NOTE: These offsets should be different from each other as maxmin
+    // might fail for highly symmetric, well-separated clusters.
     // Consider the case where the 8 clusters as a whole have octahedral
     // symmetry. In this case, R*R^T must be proprotional to the identity,
     // and eigenvalues are three-fold degenerate, because xy, yz and zx
     // plane are equivalent in terms of the maxmin optimization problem.
     // This means eigenvectors are arbitrary in this case.
-    double offset_x_ = 7.0;
-    double offset_y_ = 8.0;
-    double offset_z_ = 9.0;
+
+    // parameters for random cluster
+    const int n_grid_rand_ = 10000;
+    const int n_batch_rand_ = 100;
+    const double width_rand_ = 20.0;
+    const double xc_ = 5.0;
+    const double yc_ = 5.0;
+    const double zc_ = 7.0;
 };
 
-std::vector<double> gen_octant_cluster(int n_each, double offset_x, double offset_y, double offset_z, double width) {
+
+void gen_random(
+    int ngrid,
+    double xc,
+    double yc,
+    double zc,
+    double width,
+    std::vector<double>& grid,
+    std::vector<int>& idx
+) {
+
+    // Generates a set of points centered around (xc, yc, zc).
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_real_distribution<double> dis(-width, width);
+
+    grid.resize(3 * ngrid);
+    for (int i = 0; i < ngrid; ++i) {
+        grid[3*i    ] = xc + dis(gen);
+        grid[3*i + 1] = yc + dis(gen);
+        grid[3*i + 2] = zc + dis(gen);
+    }
+
+    idx.resize(ngrid);
+    std::iota(idx.begin(), idx.end(), 0);
+    std::shuffle(idx.begin(), idx.end(), gen);
+}
+
+
+void gen_octant(
+    int n_each,
+    double offset_x,
+    double offset_y,
+    double offset_z,
+    double width,
+    std::vector<double>& grid,
+    std::vector<int>& idx
+) {
 
     // Generates a set of points consisting of 8 well-separated, equal-sized
     // clusters located in individual octants.
 
-    std::vector<double> grid(n_each * 8 * 3);
-    int I = 0;
-
     std::random_device rd;
     std::mt19937 gen(rd());
     std::uniform_real_distribution<double> dis(-width, width);
 
+    int ngrid = 8 * n_each;
+    grid.resize(3 * ngrid);
+    int I = 0;
     for (int sign_x : {-1, 1}) {
         for (int sign_y : {-1, 1}) {
             for (int sign_z : {-1, 1}) {
-                for (int i = 0; i < n_each; ++i) {
+                for (int i = 0; i < n_each; ++i, ++I) {
                     grid[3*I    ] = sign_x * offset_x + dis(gen);
                     grid[3*I + 1] = sign_y * offset_y + dis(gen);
                     grid[3*I + 2] = sign_z * offset_z + dis(gen);
-                    ++I;
                 }
             }
         }
     }
-    return grid;
+
+    idx.resize(ngrid);
+    std::iota(idx.begin(), idx.end(), 0);
+    std::shuffle(idx.begin(), idx.end(), gen);
 }
 
+
 bool is_same_octant(int ngrid, const double* grid) {
     if (ngrid == 0) {
         return true;
     }
-    bool is_positive_x = grid[0] > 0;
-    bool is_positive_y = grid[1] > 0;
-    bool is_positive_z = grid[2] > 0;
+    const bool is_positive_x = grid[0] > 0;
+    const bool is_positive_y = grid[1] > 0;
+    const bool is_positive_z = grid[2] > 0;
     const double* end = grid + 3 * ngrid;
     for (; grid != end; grid += 3) {
         if ( is_positive_x != (grid[0] > 0) ||
@@ -80,16 +128,26 @@ bool is_same_octant(int ngrid, const double* grid) {
 }
 
 
-void BatchTest::SetUp()
+TEST_F(BatchTest, MaxMinRandomCluster)
 {
-    grid_ = gen_octant_cluster(n_each_, offset_x_, offset_y_, offset_z_, width_);
+    // This test verifies that the sizes of batches produced by maxmin
+    // do not exceed the specified limit.
 
-    idx_.resize(grid_.size() / 3);
-    std::iota(idx_.begin(), idx_.end(), 0);
+    gen_random(n_grid_rand_, xc_, yc_, zc_, width_rand_, grid_, idx_);
 
-    std::random_device rd;
-    std::mt19937 g(rd());
-    std::shuffle(idx_.begin(), idx_.end(), g);
+    std::vector<int> delim = 
+        maxmin(grid_.data(), idx_.data(), idx_.size(), n_batch_rand_);
+
+    for (size_t i = 0; i < delim.size(); ++i) {
+        if (i == 0) {
+            EXPECT_EQ(delim[i], 0);
+        } else {
+            int sz_batch = delim[i] - delim[i-1];
+            EXPECT_GT(sz_batch, 0);
+            EXPECT_LE(sz_batch, n_batch_rand_);
+        }
+    }
+    EXPECT_LE(idx_.size() - delim.back(), n_batch_rand_);
 }
 
 
@@ -99,18 +157,21 @@ TEST_F(BatchTest, MaxMinOctantCluster)
     // well-separated, equal-sized clusters located in individual octants.
     // The resulting batches should be able to recover this structure.
 
+    gen_octant(n_batch_oct_, offset_x_, offset_y_, offset_z_, width_oct_,
+               grid_, idx_);
+
     std::vector<int> delim = 
-        maxmin(grid_.data(), idx_.data(), grid_.size() / 3, n_each_);
+        maxmin(grid_.data(), idx_.data(), idx_.size(), n_batch_oct_);
 
     EXPECT_EQ(delim.size(), 8);
 
-    std::vector<double> grid_batch(3 * n_each_);
+    std::vector<double> grid_batch(3 * n_batch_oct_);
     for (int i = 0; i < 8; ++i) {
 
-        EXPECT_EQ(delim[i], i * n_each_);
+        EXPECT_EQ(delim[i], i * n_batch_oct_);
 
         // collect points within the present batch
-        for (int j = 0; j < n_each_; ++j) {
+        for (int j = 0; j < n_batch_oct_; ++j) {
             int ig = idx_[delim[i] + j];
             grid_batch[3*j    ] = grid_[3*ig    ];
             grid_batch[3*j + 1] = grid_[3*ig + 1];
@@ -118,7 +179,7 @@ TEST_F(BatchTest, MaxMinOctantCluster)
         }
 
         // verify that points in a batch reside in the same octant
-        EXPECT_TRUE(is_same_octant(n_each_, grid_batch.data()));
+        EXPECT_TRUE(is_same_octant(n_batch_oct_, grid_batch.data()));
     }
 }
 

From de983538e6ed4772297a2e788255da7f172c9a40 Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Mon, 25 Nov 2024 15:47:49 +0800
Subject: [PATCH 7/7] small fix

---
 source/module_base/grid/batch.cpp           |  2 +-
 source/module_base/grid/batch.h             |  4 +-
 source/module_base/grid/test/test_batch.cpp | 73 +++++++++++++++------
 3 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/source/module_base/grid/batch.cpp b/source/module_base/grid/batch.cpp
index 26c9f4d60e..832323e62a 100644
--- a/source/module_base/grid/batch.cpp
+++ b/source/module_base/grid/batch.cpp
@@ -35,7 +35,7 @@ namespace {
  *                          indices placed together.
  * @param[in]       m       Number of selected points (length of idx).
  *
- * @return The number of points in the first subset (according to idx).
+ * @return The number of points in the first subset within idx.
  *
  */
 int _maxmin_divide(const double* grid, int* idx, int m) {
diff --git a/source/module_base/grid/batch.h b/source/module_base/grid/batch.h
index a2646b9b70..8a6ef733e1 100644
--- a/source/module_base/grid/batch.h
+++ b/source/module_base/grid/batch.h
@@ -35,7 +35,9 @@ namespace Batch {
  *      8  9             20 10 11
  *     12 13 19           14 15
  *
- * a possible outcome with m_thr = 4 and idx(in) = {0, 1, 2, ..., 15} is:
+ * a possible outcome with m_thr = 4 and idx(in) = {0, 1, 2, ..., 15}
+ * (idx may correspond to a subset of grid and does not have to be sorted,
+ * but it must not contain duplicates) is:
  *
  * idx(out): 0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15
  * return  : {0, 4, 8, 12}
diff --git a/source/module_base/grid/test/test_batch.cpp b/source/module_base/grid/test/test_batch.cpp
index d78446db4d..74324bc47b 100644
--- a/source/module_base/grid/test/test_batch.cpp
+++ b/source/module_base/grid/test/test_batch.cpp
@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 #include <algorithm>
 #include <random>
+//#include <cstdio>
 
 #ifdef __MPI
 #include <mpi.h>
@@ -18,7 +19,7 @@ class BatchTest: public ::testing::Test
     std::vector<double> grid_;
     std::vector<int> idx_;
 
-    // parameters for octant clusters
+    // parameters for 8-octant clusters
     const int n_batch_oct_ = 10;
     const double width_oct_ = 1.0;
     const double offset_x_ = 7.0;
@@ -32,13 +33,14 @@ class BatchTest: public ::testing::Test
     // plane are equivalent in terms of the maxmin optimization problem.
     // This means eigenvectors are arbitrary in this case.
 
-    // parameters for random cluster
-    const int n_grid_rand_ = 10000;
-    const int n_batch_rand_ = 100;
-    const double width_rand_ = 20.0;
-    const double xc_ = 5.0;
-    const double yc_ = 5.0;
-    const double zc_ = 7.0;
+
+    // parameters for a random cluster
+    const int n_grid_rand_ = 1000;
+    const int n_batch_rand_ = 200;
+    const double width_rand_ = 10.0;
+    const double xc_ = 1.0;
+    const double yc_ = 1.0;
+    const double zc_ = 2.0;
 };
 
 
@@ -128,7 +130,27 @@ bool is_same_octant(int ngrid, const double* grid) {
 }
 
 
-TEST_F(BatchTest, MaxMinRandomCluster)
+bool good_batch_size(
+    const std::vector<int>& idx,
+    const std::vector<int>& delim,
+    int n_batch_thr
+) {
+    // checks if the sizes of batches are within the specified limit
+
+    bool flag = (delim[0] == 0);
+
+    size_t i = 1;
+    while (flag && i < delim.size()) {
+        int sz_batch = delim[i] - delim[i-1];
+        flag = flag && (sz_batch > 0) && (sz_batch <= n_batch_thr);
+        ++i;
+    }
+
+    return flag && ( ((int)idx.size() - delim.back()) < n_batch_thr );
+}
+
+
+TEST_F(BatchTest, MaxMinRandom)
 {
     // This test verifies that the sizes of batches produced by maxmin
     // do not exceed the specified limit.
@@ -138,20 +160,31 @@ TEST_F(BatchTest, MaxMinRandomCluster)
     std::vector<int> delim = 
         maxmin(grid_.data(), idx_.data(), idx_.size(), n_batch_rand_);
 
-    for (size_t i = 0; i < delim.size(); ++i) {
-        if (i == 0) {
-            EXPECT_EQ(delim[i], 0);
-        } else {
-            int sz_batch = delim[i] - delim[i-1];
-            EXPECT_GT(sz_batch, 0);
-            EXPECT_LE(sz_batch, n_batch_rand_);
-        }
-    }
-    EXPECT_LE(idx_.size() - delim.back(), n_batch_rand_);
+    EXPECT_TRUE(good_batch_size(idx_, delim, n_batch_rand_));
+
+    // write grid, idx & delim to file
+    //FILE* fp = fopen("grid.dat", "w");
+    //for (size_t i = 0; i < grid_.size() / 3; ++i) {
+    //    std::fprintf(fp, "% 12.6f % 12.6f % 12.6f\n",
+    //        grid_[3*i], grid_[3*i + 1], grid_[3*i + 2]);
+    //}
+    //fclose(fp);
+
+    //fp = fopen("idx.dat", "w");
+    //for (size_t i = 0; i < idx_.size(); ++i) {
+    //    std::fprintf(fp, "%d\n", idx_[i]);
+    //}
+    //fclose(fp);
+
+    //fp = fopen("delim.dat", "w");
+    //for (size_t i = 0; i < delim.size(); ++i) {
+    //    std::fprintf(fp, "%d\n", delim[i]);
+    //}
+    //fclose(fp);
 }
 
 
-TEST_F(BatchTest, MaxMinOctantCluster)
+TEST_F(BatchTest, MaxMinOctant)
 {
     // This test applies maxmin to a set of points consisting of 8
     // well-separated, equal-sized clusters located in individual octants.