Improve equality check for COO and CSR

francois141 · Oct 23, 2023 · 106b7d1 · 106b7d1
1 parent a499e9e
commit 106b7d1
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/src/matrices/COO.hpp b/src/matrices/COO.hpp
@@ -93,11 +93,26 @@ class COO {
     }
 
     friend bool operator==(const COO &lhs, const COO &rhs) {
-        return lhs.colPositions == rhs.colPositions &&
+        bool posMatch = lhs.colPositions == rhs.colPositions &&
                lhs.rowPositions == rhs.rowPositions &&
-               lhs.values == rhs.values &&
                lhs.rows == rhs.rows &&
                lhs.cols == rhs.cols;
+
+        bool valMatch = true;
+
+        if(std::is_integral<T>::value) {
+            valMatch = (lhs.values == rhs.values);
+        } else {
+            if(lhs.values.size() != rhs.values.size()) {
+                return false;
+            }
+
+            for(int i = 0; i < lhs.values.size();i++) {
+                valMatch &= abs(rhs.values[i] - lhs.values[i]) <= 1e-6;
+            }
+        }
+
+        return posMatch && valMatch;
     }
 
     int getRows() const {

diff --git a/src/matrices/CSR.hpp b/src/matrices/CSR.hpp
@@ -72,11 +72,26 @@ class CSR {
     }
 
     friend bool operator==(const CSR &lhs, const CSR &rhs) {
-        return lhs.colPositions == rhs.colPositions &&
+        bool posMatch = lhs.colPositions == rhs.colPositions &&
                lhs.rowPositions == rhs.rowPositions &&
-               lhs.values == rhs.values &&
                lhs.rows == rhs.rows &&
                lhs.cols == rhs.cols;
+
+        bool valMatch = true;
+
+        if(std::is_integral<T>::value) {
+            valMatch = (lhs.values == rhs.values);
+        } else {
+            if(lhs.values.size() != rhs.values.size()) {
+                return false;
+            }
+
+            for(int i = 0; i < lhs.values.size();i++) {
+                valMatch &= abs(rhs.values[i] - lhs.values[i]) <= 1e-6;
+            }
+        }
+
+        return posMatch && valMatch;
     }
 
     int getRows() const {