Revive DIIS Implementation

experimental/include/tensorwrapper/diis/diis.hpp → include/tensorwrapper/diis/diis.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 NWChemEx-Project
+ * Copyright 2025 NWChemEx-Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,7 +15,7 @@
  */
 
 #pragma once
-#include "tensorwrapper/tensor/tensor_wrapper.hpp"
+#include "tensorwrapper/tensor/tensor.hpp"
 #include <Eigen/Dense>
 #include <deque>
 
@@ -39,7 +39,7 @@ class DIIS {
     using size_type = std::size_t;
 
     /// Type of the value and error matrices
-    using tensor_type = tensorwrapper::tensor::ScalarTensorWrapper;
+    using tensor_type = tensorwrapper::Tensor;
 
 private:
     /// Type of the container that stores the value and error matrices
@@ -53,13 +53,14 @@ class DIIS {
       Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
 
 public:
-    /** @brief Initializes the object with the given vector size.
+    /** @brief Initializes the object with the given sample size.
      *
-     *  @param[in] max_vec The max number of previous values stored.
-     *                     Default is 5.
+     *  @param[in] max_samples The max number of previous values stored.
+     *                         Default is 5.
      */
-    DIIS(size_type max_vec = 5) :
-      m_max_vec_(max_vec), m_B_(matrix_type::Zero(max_vec, max_vec)) {}
+    DIIS(size_type max_samples = 5) :
+      m_max_samples_(max_samples),
+      m_B_(matrix_type::Zero(max_samples, max_samples)) {}
 
     /** @brief Performs DIIS extrapolation with the new value and error
      * matrices.
@@ -86,9 +87,16 @@ class DIIS {
     bool operator==(const DIIS& rhs) const noexcept;
 
 private:
-    size_type m_max_vec_;
-    deque_type m_x_values_;
+    /// The maximum number of previous samples used to extrapolate from
+    size_type m_max_samples_;
+
+    /// The queue of previous samples
+    deque_type m_samples_;
+
+    /// The queue of previous error values
     deque_type m_errors_;
+
+    /// The matrix used to solve for the mixing coefficients
     matrix_type m_B_;
 };
 

include/tensorwrapper/tensorwrapper.hpp

@@ -18,6 +18,7 @@
 #include <tensorwrapper/allocator/allocator.hpp>
 #include <tensorwrapper/buffer/buffer.hpp>
 #include <tensorwrapper/detail_/detail_.hpp>
+#include <tensorwrapper/diis/diis.hpp>
 #include <tensorwrapper/dsl/dsl.hpp>
 #include <tensorwrapper/layout/layout.hpp>
 #include <tensorwrapper/operations/operations.hpp>

experimental/src/tensorwrapper/diis/diis.cpp → src/tensorwrapper/diis/diis.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 NWChemEx-Project
+ * Copyright 2025 NWChemEx-Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,29 +14,34 @@
  * limitations under the License.
  */
 
+#include <tensorwrapper/allocator/allocator.hpp>
+#include <tensorwrapper/buffer/buffer.hpp>
 #include <tensorwrapper/diis/diis.hpp>
-#include <tensorwrapper/tensor/expression/dot.hpp>
 
 namespace tensorwrapper::diis {
 
 using tensor_type = DIIS::tensor_type;
 
 tensor_type DIIS::extrapolate(const tensor_type& X, const tensor_type& E) {
+    auto rt = X.buffer().allocator().runtime();
+    tensorwrapper::allocator::Eigen<double> allocator(rt);
+
     // Append new values to stored values
-    m_x_values_.push_back(X);
+    m_samples_.push_back(X);
     m_errors_.push_back(E);
 
     // If we're over the max number of stored values, pop the oldest ones
     // Also update m_B_ to overwrite the oldest values
-    if(m_errors_.size() > m_max_vec_) {
+    if(m_errors_.size() > m_max_samples_) {
         m_errors_.pop_front();
-        m_x_values_.pop_front();
+        m_samples_.pop_front();
 
+        // Overwrite the top-left block with the bottom right block.
         // No need to zero out the parts that aren't overwritten,
         // they'll be overwritten in the next step
-        if(m_max_vec_ > 1) {
-            m_B_.block(0, 0, m_max_vec_ - 1, m_max_vec_ - 1) =
-              m_B_.block(1, 1, m_max_vec_ - 1, m_max_vec_ - 1);
+        if(m_max_samples_ > 1) {
+            m_B_.block(0, 0, m_max_samples_ - 1, m_max_samples_ - 1) =
+              m_B_.block(1, 1, m_max_samples_ - 1, m_max_samples_ - 1);
         }
     }
 
@@ -45,11 +50,14 @@ tensor_type DIIS::extrapolate(const tensor_type& X, const tensor_type& E) {
 
     // Add the new values to m_B_
     size_type i = sz - 1;
-    for(size_type j = 0; j <= i; ++j) // compute upper triangle
-    {
+    for(size_type j = 0; j <= i; ++j) { // compute upper triangle
         tensor_type& E_i = m_errors_.at(i);
         tensor_type& E_j = m_errors_.at(j);
-        m_B_(i, j)       = tensor::expression::dot(E_i("mu,nu"), E_j("mu,nu"));
+
+        tensor_type temp;
+        temp("")               = E_i("mu,nu") * E_j("mu,nu");
+        const auto& temp_eigen = allocator.rebind(temp.buffer());
+        m_B_(i, j)             = temp_eigen.get_elem({});
 
         // Fill in lower triangle
         if(i != j) m_B_(j, i) = m_B_(i, j);
@@ -69,18 +77,18 @@ tensor_type DIIS::extrapolate(const tensor_type& X, const tensor_type& E) {
 
     // Extrapolate the new X from the coefficients.
     tensor_type new_X;
-    new_X("mu,nu") = coefs(0) * m_x_values_.at(0)("mu,nu");
-    for(int i = 1; i < sz; i++) {
+    new_X("mu,nu") = m_samples_.at(0)("mu,nu") * coefs(0);
+    for(size_type i = 1; i < sz; i++) {
         tensor_type x_i;
-        x_i("mu,nu")   = coefs(i) * m_x_values_.at(i)("mu,nu");
+        x_i("mu,nu")   = m_samples_.at(i)("mu,nu") * coefs(i);
         new_X("mu,nu") = new_X("mu,nu") + x_i("mu,nu");
     }
     return new_X;
 }
 
 bool DIIS::operator==(const DIIS& rhs) const noexcept {
-    return ((m_max_vec_ == rhs.m_max_vec_) &&
-            (m_x_values_ == rhs.m_x_values_) && (m_errors_ == rhs.m_errors_) &&
+    return ((m_max_samples_ == rhs.m_max_samples_) &&
+            (m_samples_ == rhs.m_samples_) && (m_errors_ == rhs.m_errors_) &&
             (m_B_ == rhs.m_B_));
 }
 

tests/cxx/unit_tests/tensorwrapper/diis/diis.cpp

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "../testing/testing.hpp"
+#include <tensorwrapper/diis/diis.hpp>
+#include <tensorwrapper/operations/approximately_equal.hpp>
+
+using diis_type   = tensorwrapper::diis::DIIS;
+using tensor_type = tensorwrapper::Tensor;
+using il_type     = typename tensor_type::matrix_il_type;
+
+TEST_CASE("DIIS") {
+    // Inputs
+    il_type il1{{1.0, 2.0}, {3.0, 4.0}};
+    il_type il2{{6.0, 5.0}, {8.0, 7.0}};
+    il_type il3{{12.0, 11.0}, {10.0, 9.0}};
+    tensor_type i1(il1), i2(il2), i3(il3);
+
+    SECTION("Typedefs") {
+        SECTION("size_type") {
+            using corr_t = std::size_t;
+            using the_t  = diis_type::size_type;
+            STATIC_REQUIRE(std::is_same_v<corr_t, the_t>);
+        }
+        SECTION("tensor_type") {
+            using corr_t = tensor_type;
+            using the_t  = diis_type::tensor_type;
+            STATIC_REQUIRE(std::is_same_v<corr_t, the_t>);
+        }
+    }
+
+    SECTION("Comparisons") {
+        auto defaulted        = diis_type();
+        auto two_samples_max  = diis_type(2);
+        auto extrapolate_used = diis_type();
+        auto temp             = extrapolate_used.extrapolate(i1, i3);
+        SECTION("Equals") {
+            REQUIRE(defaulted == diis_type());
+            REQUIRE(two_samples_max == diis_type(2));
+        }
+        SECTION("Max samples not equal") {
+            REQUIRE(two_samples_max != defaulted);
+        }
+        SECTION("Recorded values different") {
+            REQUIRE(defaulted != extrapolate_used);
+        }
+    }
+
+    SECTION("extrapolate") {
+        // Outputs
+        il_type il4{{12.0, 8.6}, {14.0, 10.6}};
+        il_type il5{{15.35294118, 14.35294118}, {11.11764706, 10.11764706}};
+        tensor_type corr1(il1), corr2(il4), corr3(il5);
+
+        // Call extrapolate enough to require removing an old value
+        auto diis    = diis_type(2);
+        auto output1 = diis.extrapolate(i1, i3);
+        auto output2 = diis.extrapolate(i2, i2);
+        auto output3 = diis.extrapolate(i3, i1);
+
+        using tensorwrapper::operations::approximately_equal;
+        REQUIRE(approximately_equal(output1, corr1, 1E-6));
+        REQUIRE(approximately_equal(output2, corr2, 1E-6));
+        REQUIRE(approximately_equal(output3, corr3, 1E-6));
+    }
+}