Fix missing diagonal elements in rescaled KPM Hamiltonian

This only affected models with a zero diagonal in the original
Hamiltonian matrix but with a non-zero KPM scaling factor `b`.

cppcore/src/kpm/OptimizedHamiltonian.cpp

@@ -111,13 +111,6 @@ void OptimizedHamiltonian::create_reordered(Indices const& idx, Scale<> s) {
         // corresponding to the h2_row of the reordered matrix
         auto const row = index_queue[h2_row];
         h_view.for_each_in_row(row, [&](storage_idx_t col, scalar_t value) {
-            // A diagonal element may need to be inserted into the reordered matrix
-            // even if the original matrix doesn't have an element on the main diagonal
-            if (scale.b != 0 && !diagonal_inserted && col > row) {
-                h2.insert(h2_row, h2_row) = -scale.b * inverted_a;
-                diagonal_inserted = true;
-            }
-
             // This may be a new index, map it
             if (reorder_map[col] < 0) {
                 reorder_map[col] = static_cast<storage_idx_t>(index_queue.size());
@@ -137,6 +130,12 @@ void OptimizedHamiltonian::create_reordered(Indices const& idx, Scale<> s) {
             h2.insert(h2_row, h2_col) = h2_value;
         });
 
+        // A diagonal element may need to be inserted into the reordered matrix
+        // even if the original matrix doesn't have an element on the main diagonal
+        if (scale.b != 0 && !diagonal_inserted) {
+            h2.insert(h2_row, h2_row) = -scale.b * inverted_a;
+        }
+
         // Reached the end of a slice
         if (h2_row == slice_border_indices.back() - 1) {
             slice_border_indices.push_back(static_cast<storage_idx_t>(index_queue.size()));

cppcore/tests/test_kpm.cpp

@@ -171,6 +171,18 @@ TEST_CASE("OptimizedHamiltonian reordering", "[kpm]") {
     }
 }
 
+TEST_CASE("OptimizedHamiltonian scaling") {
+    auto model = Model(graphene::monolayer(), shape::rectangle(0.6f, 0.8f));
+    auto oh = kpm::OptimizedHamiltonian(model.hamiltonian(), kpm::MatrixFormat::CSR, true);
+    auto bounds = kpm::Bounds(-12, 10); // ensures `scale.b != 0`
+    oh.optimize_for({0, 0}, bounds.scaling_factors());
+    auto scaled = oh.matrix().get<SparseMatrixX<float>>();
+
+    // Because `scale.b != 0` the scaled Hamiltonian matrix should get
+    // a non-zero diagonal even if the original matrix didn't have one.
+    REQUIRE(scaled.nonZeros() == model.hamiltonian().non_zeros() + model.hamiltonian().rows());
+}
+
 struct TestGreensResult {
     ArrayXcd g_ii, g_ij;