Fix LR optical spin weights, density restart fallback, and CSR header parsing

source/source_base/module_out/csr_reader.cpp

@@ -1,4 +1,7 @@
 #include "csr_reader.h"
+#include <cmath>
+#include <set>
+#include <limits>
 #include "source_base/tool_quit.h"
 
 namespace ModuleIO
@@ -11,180 +14,187 @@ csrFileReader<T>::csrFileReader(const std::string& filename) : FileReader(filena
     parseFile();
 }
 
-// function to parse file
+// Format-specific metadata is followed by the same CSR block reader.
 template <typename T>
 void csrFileReader<T>::parseFile()
 {
-    // Check if file is open
-    if (!isOpen())
+    RCoordinates.clear();
+    sparse_matrices.clear();
+    nextDataLine();
+    const std::string line = ss.str();
+    const std::size_t begin = line.find_first_not_of(" \t");
+    if (line.compare(begin, 5, "STEP:") == 0)
     {
-        ModuleBase::WARNING_QUIT("csrFileReader::parseFile", "File is not open");
+        parseCompactHeader();
     }
+    else
+    {
+        parseCellHeader();
+    }
+    if (matrixDimension <= 0 || numberOfR < 0)
+    {
+        ModuleBase::WARNING_QUIT("csrFileReader::parseFile", "Invalid CSR dimensions or R-block count");
+    }
+    readBlocks();
+    validateFrameEnd();
+}
 
-    std::string tmp_string;
-
-    // Read the step
-    readLine();
-    ss >> tmp_string >> tmp_string >> tmp_string >> step;
-
-    //std::cout << " step is " << step << std::endl; 
-    // Read the title
-    readLine();
-    // Read the total spin
-    readLine();
-    // Read the spin index
-    readLine();
+template <typename T>
+void csrFileReader<T>::parseCompactHeader()
+{
+    const auto invalid = []() {
+        ModuleBase::WARNING_QUIT("csrFileReader::parseCompactHeader", "Malformed compact CSR file");
+    };
+    std::string label;
+    if (!(ss >> label >> step) || label != "STEP:") { invalid(); }
+    const auto header_integer = [&](const std::string& prefix, int& value) {
+        readLine();
+        const std::string line = ss.str();
+        const std::size_t begin = line.find_first_not_of(" \t");
+        const std::size_t colon = line.find(':');
+        if (begin == std::string::npos || line.compare(begin, prefix.size(), prefix) != 0
+            || colon == std::string::npos) { invalid(); }
+        std::istringstream values(line.substr(colon + 1));
+        if (!(values >> value)) { invalid(); }
+    };
+    header_integer("Matrix Dimension of ", matrixDimension);
+    header_integer("Matrix number of ", numberOfR);
+}
 
-    // Read the matrix dimension
+// Cell-aware headers used by HContainer writers include spin and geometry.
+template <typename T>
+void csrFileReader<T>::parseCellHeader()
+{
+    const auto invalid = []() {
+        ModuleBase::WARNING_QUIT("csrFileReader::parseCellHeader", "Malformed cell-aware CSR header");
+    };
+    std::string marker;
+    std::string ionic;
+    std::string step_label;
+    if (!(ss >> marker >> ionic >> step_label >> step)
+        || marker != "---" || ionic != "Ionic" || step_label != "Step") { invalid(); }
+    readLine(); // matrix description
     readLine();
-    ss >> matrixDimension;
-    // std::cout << " mat dim is " << matrixDimension << std::endl; 
-
-    // Read the number of R
+    int spins = 0;
+    if (!(ss >> spins) || spins <= 0) { invalid(); }
     readLine();
-    ss >> numberOfR;
-
-    // Read cell
-    // Skip empty line before ucell info if present
+    int spin = 0;
+    if (!(ss >> spin) || spin < 1 || spin > spins) { invalid(); }
     readLine();
-    std::string line = ss.str();
-    if (line.empty() || line.find_first_not_of(" \t\n\r") == std::string::npos)
-    {
-        // Empty line, read next line for latName
-        readLine();
-    }
-    // Now ss contains latName (or the line after empty line)
-    // We don't need to use latName, just continue reading
-
-    // Read lat0, latvec (3 lines), atom labels (6 lines total including latName)
-    readLine(); // lat0
-    readLine(); // latvec e1
-    readLine(); // latvec e2
-    readLine(); // latvec e3
-    readLine(); // atom labels
-
-    // Read atom numbers
+    if (!(ss >> matrixDimension)) { invalid(); }
     readLine();
-    std::istringstream iss(ss.str());
-    int natom = 0;
-    int total_natom = 0;
-    while (iss >> natom)
+    if (!(ss >> numberOfR)) { invalid(); }
+    nextDataLine(); // lattice name, possibly after an empty line
+    for (int i = 0; i < 5; ++i) { readLine(); } // scale, three vectors, species
+    readLine(); // atom counts
+    int atoms = 0;
+    int count = 0;
+    while (true)
     {
-        total_natom += natom;
-    }
-
-    // Read "Direct" line and atom coordinates
-    for (int i = 0; i < total_natom + 1; i++)
-    {
-        readLine(); // Direct + atom coordinates
+        ss >> std::ws;
+        if (ss.eof() || ss.peek() == '#') { break; }
+        if (!(ss >> count) || count < 0 || count > std::numeric_limits<int>::max() - atoms)
+        { invalid(); }
+        atoms += count;
     }
+    if (atoms <= 0) { invalid(); }
+    readLine(); // coordinate convention
+    for (int i = 0; i < atoms; ++i) { readLine(); }
+}
 
-    // Skip empty lines and CSR format comment block
-    // (lines starting with '#' or containing only whitespace)
-    // Use readLine() since ifs is private in FileReader
-    bool found_data = false;
-    while (!found_data)
+// Skip only whitespace and comments, never malformed numerical data.
+template <typename T>
+void csrFileReader<T>::nextDataLine()
+{
+    do
     {
         readLine();
-        std::string line = ss.str();
-        size_t start = line.find_first_not_of(" \t");
-        if (start != std::string::npos && line[start] != '#')
-        {
-            // This is a data line (first R-block), ss already holds it
-            found_data = true;
-        }
-    }
-
-    // Read the matrices
-    // Note: ss already contains the first R-block line from the skip loop above
-    for (int i = 0; i < numberOfR; i++)
-    {
-        // std::cout << " read R " << i+1 << std::endl;
+        ss >> std::ws;
+    } while (ss.eof() || ss.peek() == '#');
+}
 
-        std::vector<int> RCoord(3);
-        int nonZero = 0;
+template <typename T>
+void csrFileReader<T>::nextDataToken()
+{
+    ss >> std::ws;
+    if (ss.eof() || ss.peek() == '#') { nextDataLine(); }
+}
 
-        if (i > 0)
+template <typename T>
+void csrFileReader<T>::readBlocks()
+{
+    const auto invalid = []() {
+        ModuleBase::WARNING_QUIT("csrFileReader::readBlocks", "Malformed CSR block");
+    };
+    std::set<std::vector<int>> seen;
+    for (int ir = 0; ir < numberOfR; ++ir)
+    {
+        nextDataLine();
+        std::vector<int> coordinate(3);
+        int nnz = 0;
+        if (!(ss >> coordinate[0] >> coordinate[1] >> coordinate[2] >> nnz) || nnz < 0)
+        { invalid(); }
+        const std::size_t dimension = matrixDimension;
+        if (static_cast<std::size_t>(nnz) > dimension * dimension || !seen.insert(coordinate).second)
+        { invalid(); }
+        std::vector<T> values(nnz);
+        std::vector<int> columns(nnz);
+        std::vector<int> pointers(dimension + 1);
+        for (int i = 0; i < nnz; ++i)
         {
-            // For subsequent R-blocks, skip empty lines to find next R-coordinate line
-            bool found = false;
-            while (!found)
-            {
-                readLine();
-                std::string line = ss.str();
-                size_t start = line.find_first_not_of(" \t");
-                if (start != std::string::npos && line[start] != '#')
-                {
-                    found = true;
-                }
-            }
+            nextDataToken();
+            if (!(ss >> values[i]) || !std::isfinite(std::abs(values[i]))) { invalid(); }
         }
-        ss >> RCoord[0] >> RCoord[1] >> RCoord[2] >> nonZero;
-        RCoordinates.push_back(RCoord);
-
-        std::vector<T> csr_values(nonZero);
-        std::vector<int> csr_col_ind(nonZero);
-        std::vector<int> csr_row_ptr(matrixDimension + 1);
-
-        // read CSR values
-        readLine();
-        // std::cout << " ss1: " << ss.str() << std::endl;
-
-        readLine();
-    size_t count1 = 0;
-        while (count1 < nonZero)
+        for (int i = 0; i < nnz; ++i)
         {
-            if (ss.eof() || ss.fail())
-            {
-                readLine();
+            nextDataToken();
+            if (!(ss >> columns[i]) || columns[i] < 0 || columns[i] >= matrixDimension) { invalid(); }
         }
-            if (ss >> csr_values[count1])
-            {
-                count1++;
-            }
-    }
-        // std::cout << "count1=" << count1 << std::endl;
-
-        // read CSR column indices
-        readLine();
-        // std::cout << " ss2: " << ss.str() << std::endl;
-
-    size_t count2 = 0;
-        while (count2 < nonZero)
+        for (int i = 0; i <= matrixDimension; ++i)
         {
-            if (ss.eof() || ss.fail())
-            {
-                readLine();
+            nextDataToken();
+            if (!(ss >> pointers[i]) || pointers[i] < 0 || pointers[i] > nnz) { invalid(); }
+            if (i > 0 && pointers[i] < pointers[i - 1]) { invalid(); }
         }
-            if (ss >> csr_col_ind[count2])
+        if (pointers.front() != 0 || pointers.back() != nnz) { invalid(); }
+        ss >> std::ws;
+        if (!ss.eof() && ss.peek() != '#') { invalid(); }
+        for (int row = 0; row < matrixDimension; ++row)
+        {
+            for (int i = pointers[row] + 1; i < pointers[row + 1]; ++i)
             {
-                count2++;
+                if (columns[i] <= columns[i - 1]) { invalid(); }
             }
+        }
+        SparseMatrix<T> matrix(matrixDimension, matrixDimension);
+        matrix.readCSR(values, columns, pointers);
+        RCoordinates.push_back(coordinate);
+        sparse_matrices.push_back(matrix);
     }
-        // std::cout << "count2=" << count2 << std::endl;
-
-        // read row pointers
-        readLine();
-        // std::cout << " ss3: " << ss.str() << std::endl;
+}
 
-    size_t count3 = 0;
-        while (count3 < matrixDimension + 1)
+// Only blank/comment lines, EOF or an unconsumed next frame may follow a frame.
+template <typename T>
+void csrFileReader<T>::validateFrameEnd()
+{
+    std::string line;
+    while (peekLine(line))
+    {
+        std::istringstream tokens(line);
+        tokens >> std::ws;
+        if (tokens.eof() || tokens.peek() == '#')
         {
-            if (ss.eof() || ss.fail())
-            {
-                readLine();
+            readLine();
+            continue;
         }
-            if (ss >> csr_row_ptr[count3])
-            {
-                count3++;
-            }
-    }
-        // std::cout << "count3=" << count3 << std::endl;
-
-        // create sparse matrix
-        SparseMatrix<T> matrix(matrixDimension, matrixDimension);
-        matrix.readCSR(csr_values, csr_col_ind, csr_row_ptr);
-        sparse_matrices.push_back(matrix);
+        std::string marker;
+        tokens >> marker;
+        if (marker == "STEP:") { return; }
+        std::string ionic;
+        std::string step_label;
+        if (marker == "---" && tokens >> ionic >> step_label
+            && ionic == "Ionic" && step_label == "Step") { return; }
+        ModuleBase::WARNING_QUIT("csrFileReader::validateFrameEnd", "Unexpected data after CSR frame");
     }
 }
 

source/source_base/module_out/csr_reader.h

@@ -12,7 +12,12 @@ namespace ModuleIO
 /**
  * @brief Class to read CSR file
  * @details This class is used to read CSR file
- * the default format is like:
+ * Both save_sparse's compact three-line header (STEP, Matrix Dimension,
+ * Matrix number) and the cell-aware format below are accepted. Compact CSR
+ * blocks contain an R coordinate/nnz line, values, zero-based column indices,
+ * and dimension+1 row pointers. Invalid dimensions, indices and truncated data
+ * are rejected for both formats rather than producing an empty matrix.
+ * The cell-aware format is like:
  * ```
  *  --- Ionic Step 1 ---
  *  # print density matrix in real space DM(R)
@@ -72,11 +77,17 @@ class csrFileReader : public FileReader
     int getMatrixDimension() const;
 
   private:
+    void parseCompactHeader();
+    void parseCellHeader();
+    void readBlocks();
+    void validateFrameEnd();
+    void nextDataLine();
+    void nextDataToken();
     std::vector<std::vector<int>> RCoordinates;
     std::vector<SparseMatrix<T>> sparse_matrices;
-    int step;
-    int matrixDimension;
-    int numberOfR;
+    int step = 0;
+    int matrixDimension = 0;
+    int numberOfR = 0;
 };
 
 } // namespace ModuleIO

source/source_base/module_out/file_reader.cpp

@@ -47,6 +47,16 @@ void FileReader::readLine()
     }
 }
 
+bool FileReader::peekLine(std::string& line)
+{
+    const std::streampos position = ifs.tellg();
+    if (position == std::streampos(-1)) { return false; }
+    const bool found = static_cast<bool>(std::getline(ifs, line));
+    ifs.clear();
+    ifs.seekg(position);
+    return found;
+}
+
 void FileReader::read_ucell()
 {
     if (ifs.eof())