From 55417a0efa573ed06e39bef53b8302717f365910 Mon Sep 17 00:00:00 2001
From: kradchen <chenjuijun@realcan.cn>
Date: Thu, 11 May 2023 13:49:58 +0800
Subject: [PATCH] Add size function

---
 src/Function3D.cpp | 29 +++++++++++++++++++++++++++++
 src/Function3D.h   |  1 +
 2 files changed, 30 insertions(+)

diff --git a/src/Function3D.cpp b/src/Function3D.cpp
index 9e93a89..cfb24e4 100644
--- a/src/Function3D.cpp
+++ b/src/Function3D.cpp
@@ -88,3 +88,32 @@ Matrix Aurora::zeros(int aRow, int aColumn, int aSlice) {
 Matrix Aurora::zeros(int aSquareRow) {
     return Aurora::zeros(aSquareRow, aSquareRow);
 }
+
+Matrix Aurora::size(const Matrix &aMatrix)
+{
+    if (aMatrix.isScalar()){
+        double * output = Aurora::malloc(1);
+        output[0]=1;
+        return Matrix::New(output,1,1,1);
+    }
+    else if (aMatrix.isVector()){
+        double * output = Aurora::malloc(1);
+        output[0]=aMatrix.getDataSize();
+        return Matrix::New(output,1,1,1);
+    }
+    //3D
+    else if (aMatrix.getDimSize(2)>1){
+        double * output = Aurora::malloc(3);
+        output[0]=aMatrix.getDimSize(0);
+        output[1]=aMatrix.getDimSize(1);
+        output[2]=aMatrix.getDimSize(2);
+        return Matrix::New(output,3,1,1);
+    }
+    //2D matrix
+    else{
+        double * output = Aurora::malloc(2);
+        output[0]=aMatrix.getDimSize(0);
+        output[1]=aMatrix.getDimSize(1);
+        return Matrix::New(output,2,1,1);
+    }
+}
diff --git a/src/Function3D.h b/src/Function3D.h
index 772cf8e..bf63753 100644
--- a/src/Function3D.h
+++ b/src/Function3D.h
@@ -41,6 +41,7 @@ namespace Aurora {
     Matrix interp3(const Matrix& aX, const Matrix& aY, const Matrix& aZ, const Matrix& aV, const Matrix& aX1, const Matrix& aY1, const Matrix& aZ1,InterpnMethod aMethod);
     Matrix interpn(const Matrix& aX, const Matrix& aY, const Matrix& aZ, const Matrix& aV, const Matrix& aX1, const Matrix& aY1, const Matrix& aZ1,InterpnMethod aMethod);
 
+    Matrix size(const Matrix &aMatrix);
 };