Reduce repeated onsite projector setup and cover mixed orbital meshes

source/source_pw/module_proj/onsite_proj_init.cpp

@@ -234,22 +234,22 @@ void projectors::OnsiteProjector<T, Device>::init_proj(const std::string& orbita
         nr = std::max(nr, nr_); // the maximal nr
         assert(dr == -1.0 || dr == dr_); // the dr should be the same for all types
         dr = (dr == -1.0) ? dr_ : dr;
+        rgrid.resize(nr_);
+        std::iota(rgrid.begin(), rgrid.end(), 0);
+        std::for_each(rgrid.begin(), rgrid.end(), [dr](double& r_i) { r_i *= dr; });
         for(int ip = 0; ip < nproj_it; ++ip)
         {
             if (!(onsite_r[idx] > 0.0) || onsite_r[idx] >= (nr_ - 1) * dr_)
             {
-                ModuleBase::WARNING_QUIT("OnsiteProjector", "onsite_radius must lie inside the orbital radial grid.");
+                ModuleBase::WARNING_QUIT("OnsiteProjector",
+                                         "onsite_radius must lie inside the orbital radial grid: " + orb_files[it]);
             }
             int l = lproj[idx];
             int izeta = iproj[idx];
             int irad = 0;
             irad = std::accumulate(nzeta.begin(), nzeta.begin() + l, irad);
             irad += izeta;
-            std::vector<double> temp = radials[irad];
-            rgrid.resize(nr_);
-            std::iota(rgrid.begin(), rgrid.end(), 0);
-            std::for_each(rgrid.begin(), rgrid.end(), [dr](double& r_i) { r_i *= dr; });
-            smoothgen(nr_, rgrid.data(), temp.data(), onsite_r[idx], projs[idx]);
+            smoothgen(nr_, rgrid.data(), radials[irad].data(), onsite_r[idx], projs[idx]);
             it2iproj[it][ip] = idx;
             ++idx;
         }

tests/17_DS_DFTU/CMakeLists.txt

@@ -6,6 +6,11 @@ enable_testing()
 # Uncomment the add_test blocks below once the feature branch is merged.
 
 find_program(BASH bash)
+find_package(Python3 REQUIRED COMPONENTS Interpreter)
+add_test(
+  NAME 17_DS_DFTU_onsite_mesh
+  COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_onsite_proj_mesh.py
+          ${ABACUS_BIN_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/50_FeO_O_first_Fe_second)
 if(ENABLE_ASAN)
   add_test(
     NAME 17_DS_DFTU_test_with_asan

tests/17_DS_DFTU/test_onsite_proj_mesh.py

@@ -0,0 +1,49 @@
+"""Check onsite projection with unequal orbital meshes and an excessive radius."""
+
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+import tempfile
+
+
+def run_case(binary, reference, radius, root):
+    case = root / str(radius)
+    case.mkdir()
+    for name in ("INPUT", "STRU", "KPT"):
+        shutil.copy2(reference / name, case / name)
+
+    orbital_dir = reference.parent.parent / "PP_ORB"
+    inp = (case / "INPUT").read_text()
+    inp = inp.replace("../../PP_ORB", str(orbital_dir))
+    inp = inp.replace("scf_nmax    100", "scf_nmax    1")
+    inp = inp.replace("onsite_radius   3.0", "onsite_radius   " + str(radius))
+    inp = "\n".join(line for line in inp.splitlines()
+                    if not line.lstrip().startswith(("dft_plus_u", "orbital_corr", "hubbard_u"))) + "\n"
+    (case / "INPUT").write_text(inp)
+
+    env = os.environ.copy()
+    env["OMP_NUM_THREADS"] = "1"
+    return subprocess.run([str(binary)], cwd=case, env=env,
+                          stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                          universal_newlines=True)
+
+
+def main():
+    binary = Path(sys.argv[1]).resolve()
+    reference = Path(sys.argv[2]).resolve()
+    with tempfile.TemporaryDirectory() as directory:
+        root = Path(directory)
+        valid = run_case(binary, reference, 3.0, root)
+        if valid.returncode != 0 or "INIT SCF" not in valid.stdout:
+            raise AssertionError("Valid onsite mesh failed:\n" + valid.stdout[-3000:])
+
+        invalid = run_case(binary, reference, 6.2, root)
+        expected = "onsite_radius must lie inside the orbital radial grid: Fe_gga_6au_100Ry_4s2p2d1f.orb"
+        if invalid.returncode == 0 or expected not in invalid.stdout:
+            raise AssertionError("Invalid onsite mesh was not rejected:\n" + invalid.stdout[-3000:])
+
+
+if __name__ == "__main__":
+    main()