#1 add func for 2d array calculations

src/main.cpp

@@ -22,6 +22,16 @@ double sum_array(const std::vector<double> &arr) {
     return sum;
 }
 
+double sum_2d_array(const std::vector<std::vector<double>> &arr) {
+    double sum = 0;
+    for (const auto &row : arr) {
+        for (double num : row) {
+            sum += num;
+        }
+    }
+    return sum;
+}
+
 
 namespace py = pybind11;
 
@@ -57,6 +67,7 @@ PYBIND11_MODULE(_core, m) {
     )pbdoc",py::arg("i"), py::arg("j"));
 
     m.def("sum_array", &sum_array, "Calculate the sum of elements in an array");
+    m.def("sum_2d_array", &sum_2d_array, "Calculate the sum of elements in a 2D array");
 
 #ifdef VERSION_INFO
     m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);

src/pybind_example/__init__.py

 from __future__ import annotations
 
-from ._core import __doc__, __version__, add, subtract, multiply, sum_array
+from ._core import (
+    __doc__,
+    __version__,
+    add,
+    subtract,
+    multiply,
+    sum_array,
+    sum_2d_array,
+)
 
-__all__ = ["__doc__", "__version__", "add", "subtract"]
+__all__ = [
+    "__doc__",
+    "__version__",
+    "add",
+    "subtract",
+    "multiply",
+    "sum_array",
+    "sum_2d_array",
+]

tests/test_basic.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import pybind_example as m
 import numpy as np
 
+
 def test_version():
     assert m.__version__ == "0.0.1"
 
@@ -15,10 +16,21 @@ def test_sub():
     assert m.subtract(1, 2) == -1
 
 
-def test_sum_array_numpy_array_1D():
-    array = np.asarray([1.,1.,2.,1.,1.])
+def test_sum_array_numpy_array_1D_float():
+    array = np.asarray([1.0, 1.0, 2.0, 1.0, 1.0])
     assert m.sum_array(array) == 6
 
-def test_sum_array_numpy_list_1D():
-    list_ = [1.,1.,2.,1.,1.]
-    assert m.sum_array(list_) == 6
\ No newline at end of file
+
+def test_sum_array_numpy_list_1D_float():
+    list_ = [1.0, 1.0, 2.0, 1.0, 1.0]
+    assert m.sum_array(list_) == 6
+
+
+def test_sum_array_numpy_array_1D_int():
+    array = np.ones(10)
+    assert m.sum_array(array) == 10
+
+
+def test_sum_array_numpy_array_2D_float():
+    array = np.arange(9).reshape((3,3))
+    assert m.sum_2d_array(array) == np.sum(array)