Feat SimpleArray.hpp, kalman_filter.hpp and test_gemm.py: Modify matmul() to support the argument of 1 dimensional and 2 dimensional SimpleArray

Author: ThreeMonth03

cpp/modmesh/buffer/SimpleArray.hpp

@@ -937,8 +937,8 @@ detail::SimpleArrayMixinCalculators<A, T>::median_freq(small_vector<value_type>
 }
 
 /**
- * Perform matrix multiplication for 2D arrays.
- * This implementation supports only 2D x 2D matrix multiplication.
+ * Perform matrix multiplication for SimpleArrays.
+ * This implementation supports 1D x 1D, 1D x 2D, 2D x 1D, and 2D x 2D matrix multiplication.
  */
 template <typename A, typename T>
 A SimpleArrayMixinCalculators<A, T>::matmul(A const & other) const
@@ -969,43 +969,112 @@ A SimpleArrayMixinCalculators<A, T>::matmul(A const & other) const
         }
     };
 
-    if (this_ndim != 2 || other_ndim != 2)
+    auto check_product_shape = [&](A const * athis, A const * other, ssize_t athis_idx, ssize_t other_idx) -> void
+    {
+        if (athis->shape(athis_idx) != other->shape(other_idx))
+        {
+            throw std::out_of_range(
+                std::format("SimpleArray::matmul(): shape mismatch: this={} other={}",
+                            format_shape(athis),
+                            format_shape(other)));
+        }
+    };
+
+    if ((this_ndim != 2 && this_ndim != 1) || (other_ndim != 2 && other_ndim != 1))
     {
         const std::string err = std::format("SimpleArray::matmul(): unsupported dimensions: "
-                                            "this={} other={}. Only 2D x 2D matrix multiplication is supported",
+                                            "this={} other={}. SimpleArray must be 1D or 2D.",
                                             format_shape(athis),
                                             format_shape(&other));
         throw std::out_of_range(err);
     }
 
-    const size_t m = athis->shape(0);
-    const size_t k = athis->shape(1);
-    const size_t n = other.shape(1);
+    bool this_is_1d = (this_ndim == 1);
+    bool other_is_1d = (other_ndim == 1);
 
-    if (k != other.shape(0))
+    // 1D x 1D
+    if (this_is_1d && other_is_1d)
     {
-        throw std::out_of_range(
-            std::format("SimpleArray::matmul(): shape mismatch: this={} other={}",
-                        format_shape(athis),
-                        format_shape(&other)));
+        check_product_shape(athis, &other, 0, 0);
+        A result(1);
+        value_type v = 0;
+        for (size_t i = 0; i < athis->shape(0); ++i)
+        {
+            v += (*athis)(i)*other.data(i);
+        }
+        result.data(0) = v;
+        return result;
     }
-
-    typename detail::SimpleArrayInternalTypes<T>::shape_type result_shape{m, n};
-    A result(result_shape);
-    result.fill(static_cast<value_type>(0));
-
-    for (size_t i = 0; i < m; ++i)
+    // 1D x 2D
+    else if (this_is_1d)
     {
+        const size_t k = athis->shape(0);
+        const size_t n = other.shape(1);
+        check_product_shape(athis, &other, 0, 0);
+        A result(n);
+
         for (size_t j = 0; j < n; ++j)
         {
+            value_type v = 0;
+            for (size_t l = 0; l < k; ++l)
+            {
+                v += (*athis)(l)*other(l, j);
+            }
+            result.data(j) = v;
+        }
+        return result;
+    }
+    // 2D x 1D
+    else if (other_is_1d)
+    {
+        const size_t m = athis->shape(0);
+        const size_t k = athis->shape(1);
+
+        check_product_shape(athis, &other, 1, 0);
+        A result(m);
+
+        for (size_t i = 0; i < m; ++i)
+        {
+            value_type v = 0;
             for (size_t l = 0; l < k; ++l)
             {
-                result(i, j) += athis->operator()(i, l) * other(l, j);
+                v += (*athis)(i, l) * other(l);
+            }
+            result.data(i) = v;
+        }
+        return result;
+    }
+    // 2D x 2D
+    else
+    {
+        const size_t m = athis->shape(0);
+        const size_t k = athis->shape(1);
+        const size_t n = other.shape(1);
+        check_product_shape(athis, &other, 1, 0);
+
+        shape_type result_shape{m, n};
+        A result(result_shape);
+
+        for (size_t i = 0; i < m; ++i)
+        {
+            for (size_t j = 0; j < n; ++j)
+            {
+                value_type v = 0;
+                for (size_t l = 0; l < k; ++l)
+                {
+                    v += (*athis)(i, l) * other(l, j);
+                }
+                result(i, j) = v;
             }
         }
+        return result;
     }
 
-    return result;
+    throw std::out_of_range(
+        std::format("SimpleArray::matmul(): this={} other={}"
+                    " cannot perform matrix multiplication.",
+                    format_shape(athis),
+                    format_shape(&other)));
 }
 
 /**

cpp/modmesh/linalg/kalman_filter.hpp

@@ -367,19 +367,14 @@ template <typename T>
 void KalmanFilter<T>::predict_state()
 {
     // x <- F x  (m_x <- m_f @ m_x)
-    array_type x_col = m_x.reshape(small_vector<size_t>{m_state_size, 1});
-    x_col = m_f.matmul(x_col);
-    m_x = x_col.reshape(small_vector<size_t>{m_state_size});
+    m_x = m_f.matmul(m_x);
 }
 
 template <typename T>
 void KalmanFilter<T>::predict_state(array_type const & u)
 {
     // x <- F x + B u  (m_x <- m_f @ m_x + m_b @ u)
-    array_type x_col = m_x.reshape(small_vector<size_t>{m_state_size, 1});
-    array_type u_col = u.reshape(small_vector<size_t>{m_control_size, 1});
-    x_col = m_f.matmul(x_col).add(m_b.matmul(u_col));
-    m_x = x_col.reshape(small_vector<size_t>{m_state_size});
+    m_x = m_f.matmul(m_x).add(m_b.matmul(u));
 }
 
 template <typename T>
@@ -440,10 +435,7 @@ template <typename T>
 typename KalmanFilter<T>::array_type KalmanFilter<T>::innovation(array_type const & z)
 {
     // y <- z - H x  (y <- z - m_h @ m_x)
-    array_type x_col = m_x.reshape(small_vector<size_t>{m_state_size, 1});
-    array_type hx = m_h.matmul(x_col);
-    array_type hx_vec = hx.reshape(small_vector<size_t>{m_measurement_size});
-    return z.sub(hx_vec);
+    return z.sub(m_h.matmul(m_x));
 }
 
 template <typename T>
@@ -471,10 +463,7 @@ template <typename T>
 void KalmanFilter<T>::update_state(array_type const & k, array_type const & y)
 {
     // x <- x + K y  (m_x <- m_x + k @ y)
-    array_type y_col = y.reshape(small_vector<size_t>{m_measurement_size, 1});
-    array_type ky = k.matmul(y_col);
-    array_type ky_vec = ky.reshape(small_vector<size_t>{m_state_size});
-    m_x = m_x.add(ky_vec);
+    m_x = m_x.add(k.matmul(y));
 }
 
 template <typename T>

tests/test_gemm.py

@@ -207,10 +207,40 @@ def test_compare_with_numpy(self):
             [8.0, 11.0, 11.0]
         ], dtype=self.dtype)
 
+        # Test case 4: (4x6) x (6)
+        a_data_4 = np.array([
+            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
+            [7.0, 8.0, 9.0, 10.0, 11.0, 12.0],
+            [13.0, 14.0, 15.0, 16.0, 17.0, 18.0],
+            [19.0, 20.0, 21.0, 22.0, 23.0, 24.0]
+        ], dtype=self.dtype)
+        b_data_4 = np.array([1., 2., 3., 4., 5., 6], dtype=self.dtype)
+        expected_4 = np.array([91., 217., 343., 469.], dtype=self.dtype)
+
+        # Test case 5: (6) x (6 x 4)
+        a_data_5 = np.array([1., 2., 3., 4., 5., 6], dtype=self.dtype)
+        b_data_5 = np.array([
+            [1.0, 2.0, 3.0, 4.0],
+            [5.0, 6.0, 7.0, 8.0],
+            [9.0, 10.0, 11.0, 12.0],
+            [13.0, 14.0, 15.0, 16.0],
+            [17.0, 18.0, 19.0, 20.0],
+            [21.0, 22.0, 23.0, 24.0]
+        ], dtype=self.dtype)
+        expected_5 = np.array([301., 322., 343., 364.], dtype=self.dtype)
+
+        # Test case 6: (3) x (3)
+        a_data_6 = np.array([1., 2., 3.], dtype=self.dtype)
+        b_data_6 = np.array([4., 5., 6.], dtype=self.dtype)
+        expected_6 = np.array([32.], dtype=self.dtype)
+
         test_cases = [
             (a_data_1, b_data_1, expected_1, "2x3 x 3x4"),
             (a_data_2, b_data_2, expected_2, "4x6 x 6x3"),
-            (a_data_3, b_data_3, expected_3, "3x3 x 3x3")
+            (a_data_3, b_data_3, expected_3, "3x3 x 3x3"),
+            (a_data_4, b_data_4, expected_4, "4x6 x 6"),
+            (a_data_5, b_data_5, expected_5, "6 x 6x3"),
+            (a_data_6, b_data_6, expected_6, "3 x 3x3")
         ]
 
         for a_data, b_data, expected, description in test_cases:
@@ -234,64 +264,66 @@ def test_compare_with_numpy(self):
                     np.testing.assert_array_almost_equal(
                         result.ndarray, expected, decimal=10)
 
-    def test_unsupported_dimensions_error(self):
-        """Test error handling for unsupported dimensions"""
+    def test_wrong_shape_error(self):
+        """Test error handling for wrong shapes"""
 
-        # Test 1D x 1D (not supported)
-        a_1d = self.SimpleArray(array=np.array([1.0, 2.0, 3.0],
-                                               dtype=self.dtype))
-        b_1d = self.SimpleArray(array=np.array([4.0, 5.0, 6.0],
-                                               dtype=self.dtype))
+        a_3d_data = np.array([[[1.0, 2.0], [3.0, 4.0]],
+                              [[5.0, 6.0], [7.0, 8.0]]], dtype=self.dtype)
+        b_3d_data = np.array([[[1.0, 0.0], [0.0, 1.0]],
+                              [[2.0, 0.0], [0.0, 2.0]]], dtype=self.dtype)
+        a_3d = self.SimpleArray(array=a_3d_data)
+        b_3d = self.SimpleArray(array=b_3d_data)
 
         with self.assertRaisesRegex(
             IndexError,
-            r"SimpleArray::matmul\(\): unsupported dimensions: this=\(3\) "
-            r"other=\(3\)\. Only 2D x 2D matrix multiplication is supported"
+            r"SimpleArray::matmul\(\): unsupported dimensions: "
+            r"this=\(2,2,2\) other=\(2,2,2\)\. SimpleArray must be 1D or 2D."
         ):
-            a_1d.matmul(b_1d)
-
-        # Test 1D x 2D (not supported)
-        a_1d = self.SimpleArray(array=np.array([1.0, 2.0], dtype=self.dtype))
-        b_2d = self.SimpleArray(array=np.array([[1.0, 2.0], [3.0, 4.0]],
-                                               dtype=self.dtype))
+            a_3d.matmul(b_3d)
 
+        a = np.zeros((3, 3), dtype=self.dtype)
+        b = np.zeros((2, 3), dtype=self.dtype)
+        a = self.SimpleArray(array=a)
+        b = self.SimpleArray(array=b)
         with self.assertRaisesRegex(
             IndexError,
-            r"SimpleArray::matmul\(\): unsupported dimensions: this=\(2\) "
-            r"other=\(2,2\)\. Only 2D x 2D matrix multiplication is supported"
+            r"SimpleArray::matmul\(\): shape mismatch: "
+            r"this=\(3,3\) other=\(2,3\)"
         ):
-            a_1d.matmul(b_2d)
-
-        # Test 2D x 1D (not supported)
-        a_2d = self.SimpleArray(array=np.array([[1.0, 2.0, 3.0],
-                                               [4.0, 5.0, 6.0]],
-                                               dtype=self.dtype))
-        b_1d = self.SimpleArray(array=np.array([7.0, 8.0, 9.0],
-                                               dtype=self.dtype))
+            a.matmul(b)
 
+        a = np.zeros((3, 3), dtype=self.dtype)
+        b = np.zeros((2), dtype=self.dtype)
+        a = self.SimpleArray(array=a)
+        b = self.SimpleArray(array=b)
         with self.assertRaisesRegex(
             IndexError,
-            r"SimpleArray::matmul\(\): unsupported dimensions: this=\(2,3\) "
-            r"other=\(3\)\. Only 2D x 2D matrix multiplication is supported"
+            r"SimpleArray::matmul\(\): shape mismatch: "
+            r"this=\(3,3\) other=\(2\)"
         ):
-            a_2d.matmul(b_1d)
-
-        # Test 3D x 3D (not supported - tensor operation)
-        a_3d_data = np.array([[[1.0, 2.0], [3.0, 4.0]],
-                              [[5.0, 6.0], [7.0, 8.0]]], dtype=self.dtype)
-        b_3d_data = np.array([[[1.0, 0.0], [0.0, 1.0]],
-                              [[2.0, 0.0], [0.0, 2.0]]], dtype=self.dtype)
+            a.matmul(b)
 
-        a_3d = self.SimpleArray(array=a_3d_data)
-        b_3d = self.SimpleArray(array=b_3d_data)
+        a = np.zeros((2), dtype=self.dtype)
+        b = np.zeros((3, 3), dtype=self.dtype)
+        a = self.SimpleArray(array=a)
+        b = self.SimpleArray(array=b)
+        with self.assertRaisesRegex(
+            IndexError,
+            r"SimpleArray::matmul\(\): shape mismatch: "
+            r"this=\(2\) other=\(3,3\)"
+        ):
+            a.matmul(b)
 
+        a = np.zeros((2), dtype=self.dtype)
+        b = np.zeros((3), dtype=self.dtype)
+        a = self.SimpleArray(array=a)
+        b = self.SimpleArray(array=b)
         with self.assertRaisesRegex(
             IndexError,
-            r"SimpleArray::matmul\(\): unsupported dimensions: "
-            r"this=\(2,2,2\) other=\(2,2,2\)\. Only 2D x 2D matrix "
-            r"multiplication is supported"
+            r"SimpleArray::matmul\(\): shape mismatch: "
+            r"this=\(2\) other=\(3\)"
         ):
-            a_3d.matmul(b_3d)
+            a.matmul(b)
 
     def test_matmul_operator(self):
         """Test @ operator for matrix multiplication"""