From 1a29baeef83836c084041dcbe6461de831be5c83 Mon Sep 17 00:00:00 2001
From: kradchen <chenjuijun@realcan.cn>
Date: Thu, 8 Jun 2023 15:59:57 +0800
Subject: [PATCH] Add compareSet to 1D

---
 src/Function1D.cpp       | 24 ++++++++++++++++++++++++
 src/Function1D.h         |  4 ++++
 test/Function1D_Test.cpp | 15 +++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/src/Function1D.cpp b/src/Function1D.cpp
index af170f6..c5aad1d 100644
--- a/src/Function1D.cpp
+++ b/src/Function1D.cpp
@@ -904,6 +904,30 @@ void Aurora::padding(Matrix &aMatrix, int aIndex, double aValue)
     aMatrix = Matrix::New(newData,size,1,1,aMatrix.getValueType());
 }
 
+void Aurora::compareSet(Matrix& aMatrix,double compareValue, double newValue,CompareOp op){
+    Eigen::Map<Eigen::VectorXd> v(aMatrix.getData(),aMatrix.getDataSize());
+    switch (op) {
+        case EQ:
+            v = (v.array() == compareValue).select(newValue, v);
+            break;
+        case GT:
+            v = (v.array() > compareValue).select(newValue, v);
+            break;
+        case LT:
+            v = (v.array() < compareValue).select(newValue, v);
+            break;
+        case NG:
+            v = (v.array() <= compareValue).select(newValue, v);
+            break;
+        case NL:
+            v = (v.array() >= compareValue).select(newValue, v);
+            break;
+        case NE:
+            v = (v.array() != compareValue).select(newValue, v);
+            break;
+    }
+}
+
 Matrix Aurora::convertfp16tofloat(short* aData, int aRows, int aColumns)
 {
 
diff --git a/src/Function1D.h b/src/Function1D.h
index 4a12c1e..e89f9b3 100644
--- a/src/Function1D.h
+++ b/src/Function1D.h
@@ -131,6 +131,10 @@ namespace Aurora {
 
     Matrix auroraNot(const Matrix& aMatrix);
     Matrix auroraNot(Matrix&& aMatrix);
+    enum CompareOp{
+        EQ,GT,LT,NG,NL,NE
+    };
+    void compareSet(Matrix& aMatrix,double compareValue, double newValue,CompareOp op);
     
     Matrix convertfp16tofloat(short* aData, int aRows, int aColumns);
 };
diff --git a/test/Function1D_Test.cpp b/test/Function1D_Test.cpp
index 84f1339..8b4a114 100644
--- a/test/Function1D_Test.cpp
+++ b/test/Function1D_Test.cpp
@@ -26,6 +26,21 @@ protected:
     }
 
 };
+TEST_F(Function1D_Test,compareSet){
+    double * dataA =Aurora::malloc(9);
+    for (int i = 0; i < 9; ++i) {
+        dataA[i]=(double)(i-3);
+    }
+    Aurora::Matrix A = Aurora::Matrix::New(dataA,3,3);
+    EXPECT_EQ(-3, A[0]);
+    EXPECT_EQ(-2, A[1]);
+    EXPECT_EQ(-1, A[2]);
+    compareSet(A,0,0,Aurora::LT);
+    EXPECT_EQ(0, A[0]);
+    EXPECT_EQ(0, A[1]);
+    EXPECT_EQ(0, A[2]);
+}
+
 
 TEST_F(Function1D_Test, sign) {
     double * dataA =Aurora::malloc(9);