From cbbcca5ceb1d8e14ecc604f7ba51d3c7ca7da2e6 Mon Sep 17 00:00:00 2001
From: kradchen <chenjuijun@realcan.cn>
Date: Wed, 10 May 2023 13:46:40 +0800
Subject: [PATCH] Add sub2ind and test

---
 src/Function2D.cpp       | 35 +++++++++++++++++++++++++++++++++++
 src/Function2D.h         |  6 ++++++
 test/Function1D_Test.cpp |  2 +-
 test/Function2D_Test.cpp | 17 +++++++++++++++++
 4 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/src/Function2D.cpp b/src/Function2D.cpp
index c7ba07c..436dd1d 100644
--- a/src/Function2D.cpp
+++ b/src/Function2D.cpp
@@ -1,9 +1,12 @@
+#include <algorithm>
+#include <cstddef>
 #include <iostream>
 #include "Function.h"
 #include "Function2D.h"
 #include "Function1D.h"
 //必须在Eigen之前
 #include "AuroraDefs.h"
+#include "Matrix.h"
 
 #include <Eigen/Core>
 #include <Eigen/Eigen>
@@ -813,3 +816,35 @@ Matrix Aurora::dot(const Matrix &aMatrix,const Matrix& aOther,FunctionDirection
     }
 }
 
+
+Matrix Aurora::sub2ind(const Matrix &aVMatrixSize, std::initializer_list<Matrix> aSliceIdxsList)
+{
+    if (aSliceIdxsList.size() != aVMatrixSize.getDataSize()){
+        std::cerr<<"sub2ind size not match"<<std::endl;
+        return Matrix();
+    }
+    if (aSliceIdxsList.size() == 0){
+        std::cerr<<"sub2ind no index need calc!"<<std::endl;
+        return Matrix();
+    }
+    size_t returnVectorSize = aSliceIdxsList.begin()->getDataSize();
+    double *strides =new double[aVMatrixSize.getDataSize()+1];
+    strides[0] = 1;
+    for (int i = 1; i<=aVMatrixSize.getDataSize(); ++i) {
+        strides[i] = strides[i-1]*aVMatrixSize.getData()[i-1];
+    } 
+    double* output = Aurora::malloc(returnVectorSize);
+    for (size_t i = 0; i<returnVectorSize; ++i) {
+        size_t j = 0;
+        std::for_each(aSliceIdxsList.begin(), aSliceIdxsList.end(), [strides,output,i,&aVMatrixSize,&j](const Matrix& matrix){
+            if(j == 0 ){
+                output[i]= (matrix.getData()[i])*strides[j];
+            }
+            else{
+                output[i]+= (matrix.getData()[i]-1)*strides[j];
+            }
+            ++j;
+        });
+    }
+    return Matrix::New(output,returnVectorSize,1,1);
+}
diff --git a/src/Function2D.h b/src/Function2D.h
index f4bc471..3a3bbec 100644
--- a/src/Function2D.h
+++ b/src/Function2D.h
@@ -1,6 +1,9 @@
 #ifndef AURORA_FUNCTION2D_H
 #define AURORA_FUNCTION2D_H
 
+#include <cstdio>
+#include <initializer_list>
+
 #include "Matrix.h"
 #include "Function1D.h"
 
@@ -144,6 +147,9 @@ namespace Aurora
     Matrix prod(const Matrix &aMatrix);
 
     Matrix dot(const Matrix &aMatrix, const Matrix &aOther, FunctionDirection direction = Column);
+
+    Matrix sub2ind(const Matrix &aVMatrixSize, std::initializer_list<Matrix> aSliceIdxs);
+    
 };
 
 #endif // AURORA_FUNCTION2D_H
diff --git a/test/Function1D_Test.cpp b/test/Function1D_Test.cpp
index a8b51f2..3aca3b7 100644
--- a/test/Function1D_Test.cpp
+++ b/test/Function1D_Test.cpp
@@ -118,7 +118,7 @@ TEST_F(Function1D_Test, polyval){
     auto matrixP = Aurora::Matrix::fromRawData(dataP,6);
     auto matrixX = Aurora::Matrix::fromRawData(dataX,1);
     auto result = Aurora::polyval(matrixP,matrixX);
-    EXPECT_DOUBLE_EQ(1495.6, result.getData()[0]);
+    EXPECT_DOUBLE_EQ(1495.5717310876246, result.getData()[0]);
     // EXPECT_DOUBLE_EQ(162., result.getData()[1]);
     // EXPECT_DOUBLE_EQ(262., result.getData()[2]);
 }
diff --git a/test/Function2D_Test.cpp b/test/Function2D_Test.cpp
index 2b0bc00..f232c41 100644
--- a/test/Function2D_Test.cpp
+++ b/test/Function2D_Test.cpp
@@ -480,3 +480,20 @@ TEST_F(Function2D_Test, interp2) {
     EXPECT_DOUBLE_AE(result.getData()[2],18.6089);
 }
 
+TEST_F(Function2D_Test, sub2ind) {
+    double* dI1= new double[4]{1,2,1,2};
+    Aurora::Matrix I1(std::shared_ptr<double>(dI1,std::default_delete<double[]>()),std::vector<int>{4});
+    double* dI2= new double[4]{2,2,1,1};
+    Aurora::Matrix I2(std::shared_ptr<double>(dI2,std::default_delete<double[]>()),std::vector<int>{4});
+    double* dI4= new double[4]{1,1,2,2};
+    Aurora::Matrix I3(std::shared_ptr<double>(dI4,std::default_delete<double[]>()),std::vector<int>{4});
+    double* dsz= new double[3]{2,2,2};
+    Aurora::Matrix sz(std::shared_ptr<double>(dsz,std::default_delete<double[]>()),std::vector<int>{3});
+    auto ma = Aurora::sub2ind(sz, {I1, I2, I3});
+    EXPECT_DOUBLE_EQ(3, ma.getData()[0]);
+    EXPECT_DOUBLE_EQ(4, ma.getData()[1]);
+    EXPECT_DOUBLE_EQ(5, ma.getData()[2]);
+    EXPECT_DOUBLE_EQ(6, ma.getData()[3]);
+
+}
+