diff --git a/src/main.cpp b/src/main.cpp index 9e3a86fceac139c84e6311e6a526eba9398d51a6..fac6b555d8aba64f86ba516cdb37abbd2a628eeb 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -32,6 +32,40 @@ double sum_2d_array(const std::vector<std::vector<double>> &arr) { return sum; } +std::vector<double> sum_rows(const std::vector<std::vector<double>> &arr) { + std::vector<double> row_sums; + for (const auto &row : arr) { + double row_sum = 0; + for (double num : row) { + row_sum += num; + } + row_sums.push_back(row_sum); + } + return row_sums; +} + +std::vector<std::vector<double>> matrix_multiply(const std::vector<std::vector<double>> &arr1, const std::vector<std::vector<double>> &arr2) { + int m1 = arr1.size(); + int n1 = arr1[0].size(); + int m2 = arr2.size(); + int n2 = arr2[0].size(); + + if (n1 != m2) { + throw std::invalid_argument("Matrices cannot be multiplied due to incompatible dimensions"); + } + std::vector<std::vector<double>> result(m1, std::vector<double>(n2, 0)); + + for (int i = 0; i < m1; ++i) { + for (int j = 0; j < n2; ++j) { + for (int k = 0; k < n1; ++k) { + result[i][j] += arr1[i][k] * arr2[k][j]; + } + } + } + + return result; +} + namespace py = pybind11; @@ -68,6 +102,8 @@ PYBIND11_MODULE(_core, m) { 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"); + m.def("sum_rows", &sum_rows, "Calculate the sum of elements in each row of a 2D array"); + m.def("matrix_multiply", &matrix_multiply, "Perform matrix multiplication of two 2D arrays"); #ifdef VERSION_INFO m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO); diff --git a/src/pybind_example/__init__.py b/src/pybind_example/__init__.py index bd58be87324048ec5673ec83182e809be2cd3069..a997bf6fb11f8f8aba4e76367617072d0a9a4c18 100644 --- a/src/pybind_example/__init__.py +++ b/src/pybind_example/__init__.py @@ -8,6 +8,8 @@ from ._core import ( multiply, sum_array, sum_2d_array, + sum_rows, + matrix_multiply, ) __all__ = [ @@ -18,4 +20,6 @@ __all__ = [ "multiply", "sum_array", "sum_2d_array", + "sum_rows", + "matrix_multiply", ] diff --git a/tests/test_basic.py b/tests/test_basic.py index 55efadf1d3241db6315662269bda4af4491dc419..fea12f4c8430a0158f2248aec0075076ccae8311 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -2,7 +2,7 @@ from __future__ import annotations import pybind_example as m import numpy as np - +import pytest def test_version(): assert m.__version__ == "0.0.1" @@ -32,5 +32,26 @@ def test_sum_array_numpy_array_1D_int(): def test_sum_array_numpy_array_2D_float(): - array = np.arange(9).reshape((3,3)) + array = np.arange(9).reshape((3, 3)) assert m.sum_2d_array(array) == np.sum(array) + + +def test_sum_rows_numpy_array(): + array = np.arange(9).reshape((3, 3)) + np.testing.assert_array_almost_equal(m.sum_rows(array), np.sum(array, axis=1)) + + +def test_matrix_multiply_1(): + array_1 = np.asarray([[1, 2], [3, 4]]) + array_2 = np.asarray([[5, 6], [7, 8]]) + np.testing.assert_array_almost_equal( + m.matrix_multiply(array_1, array_2), array_1 @ array_2 + ) + +def test_matrix_multiply_raise(): + array_1 = np.asarray([[1, 2, 6], [3, 4, 5]]) + array_2 = np.asarray([[5, 6, 1], [7, 8, 1]]) + with pytest.raises(ValueError) as e: + m.matrix_multiply(array_1, array_2) + assert str(e) == "Matrices cannot be multiplied due to incompatible dimensions" +