First matmul implementation

doc/reference/array_operations.rst

@@ -6,3 +6,4 @@ Operations with arrays
 
     lazy_functions
     reduction_functions
+    linear_algebra

doc/reference/linear_algebra.rst

@@ -0,0 +1,14 @@
+.. _linear_algebra:
+
+Linear Algebra
+--------------
+
+The next functions can be used for computing linear algebra operations with :ref:`NDArray <NDArray>`.
+
+.. currentmodule:: blosc2
+
+.. autosummary::
+   :toctree: autofiles/operations_with_arrays/
+   :nosignatures:
+
+    matmul

src/blosc2/__init__.py

@@ -236,6 +236,7 @@ class Tuner(Enum):
     ones,
     full,
     save,
+    matmul,
 )
 
 from .c2array import c2context, C2Array, URLPath

src/blosc2/ndarray.py

@@ -1845,14 +1845,14 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
 
         return ndslice
 
-    def squeeze(self) -> None:
+    def squeeze(self) -> NDArray:
         """Remove single-dimensional entries from the shape of the array.
 
         This method modifies the array in-place, removing any dimensions with size 1.
 
         Returns
         -------
-        out: None
+        out: NDArray
 
         Examples
         --------
@@ -1868,6 +1868,7 @@ def squeeze(self) -> None:
         (23, 11)
         """
         super().squeeze()
+        return self
 
     def indices(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
         """
@@ -3643,6 +3644,110 @@ def sort(array: NDArray, order: str | list[str] | None = None, **kwargs: Any) ->
     return larr.sort(order).compute(**kwargs)
 
 
+def matmul(x1: NDArray, x2: NDArray, **kwargs: Any) -> NDArray:
+    """
+    Computes the matrix product between two Blosc2 NDArrays.
+
+    Parameters
+    ----------
+    x1: `NDArray`
+        The first input array.
+    x2: `NDArray`
+        The second input array.
+    kwargs: Any, optional
+        Keyword arguments that are supported by the :func:`empty` constructor.
+
+    Returns
+    -------
+    out: :ref:`NDArray`
+        The matrix product of the inputs. This is a scalar only when both x1,
+        x2 are 1-d vectors.
+
+    Raises
+    ------
+    ValueError
+        If the last dimension of ``x1`` is not the same size as
+        the second-to-last dimension of ``x2``.
+
+        If a scalar value is passed in.
+
+    References
+    ----------
+    `numpy.matmul <https://numpy.org/doc/stable/reference/generated/numpy.matmul.html>`_
+
+    Examples
+    --------
+    For 2-D arrays it is the matrix product:
+
+    >>> import numpy as np
+    >>> import blosc2
+    >>> a = np.array([[1, 2],
+    ...               [3, 4]])
+    >>> nd_a = blosc2.asarray(a)
+    >>> b = np.array([[2, 3],
+    ...               [2, 1]])
+    >>> nd_b = blosc2.asarray(b)
+    >>> blosc2.matmul(nd_a, nd_b)
+    array([[ 6,  5],
+           [14, 13]])
+
+
+    For 2-D mixed with 1-D, the result is the usual.
+
+    >>> a = np.array([[1, 3],
+    ...               [0, 1]])
+    >>> nd_a = blosc2.asarray(a)
+    >>> v = np.array([1, 2])
+    >>> nd_v = blosc2.asarray(v)
+    >>> blosc2.matmul(nd_a, nd_v)
+    array([7, 2])
+    >>> blosc2.matmul(nd_v, nd_a)
+    array([1, 5])
+
+    """
+
+    # Validate arguments are not scalars
+    if np.isscalar(x1) or np.isscalar(x2):
+        raise ValueError("Arguments can't be scalars.")
+
+    # Promote 1D arrays to 2D if necessary
+    x1_is_vector = False
+    x2_is_vector = False
+    if x1.ndim == 1:
+        x1 = x1.reshape((1, x1.shape[0]))  # (N,) -> (1, N)
+        x1_is_vector = True
+    if x2.ndim == 1:
+        x2 = x2.reshape((x2.shape[0], 1))  # (M,) -> (M, 1)
+        x2_is_vector = True
+
+    # Validate matrix multiplication compatibility
+    if x1.shape[-1] != x2.shape[-2]:
+        raise ValueError("Shapes are not aligned for matrix multiplication.")
+
+    n, l = x1.shape[-2:]
+    m = x2.shape[-1]
+
+    p1, q1 = x1.chunks[-2:]
+    q2 = x2.chunks[-1]
+
+    result = blosc2.zeros((n, m), dtype=x1.dtype)
+
+    for row in range(0, n, p1):
+        row_end = (row+p1) if (row+p1) < n else n
+        for col in range(0, m, q2):
+            col_end = (col+q2) if (col+q2) < m else m
+            for aux in range(0, l, q1):
+                aux_end = (aux+q1) if (aux+q1) < l else l
+                bx1 = x1[row:row_end, aux:aux_end]
+                bx2 = x2[aux:aux_end, col:col_end]
+                result[row:row_end, col:col_end] += np.matmul(bx1, bx2)
+
+    if x1_is_vector and x2_is_vector:
+        return result[0][0]
+
+    return result.squeeze()
+
+
 # Class for dealing with fields in an NDArray
 # This will allow to access fields by name in the dtype of the NDArray
 class NDField(Operand):

tests/ndarray/test_matmul.py

@@ -0,0 +1,84 @@
+import pytest
+import numpy as np
+import blosc2
+
+
+@pytest.mark.parametrize(
+    ("ashape", "achunks", "ablocks"),
+    [
+        ((12, 10), (7, 5), (3, 3)),
+        ((10,), (9,), (7,)),
+    ],
+)
+@pytest.mark.parametrize(
+    ("bshape", "bchunks", "bblocks"),
+    [
+        ((10,), (4,), (2,)),
+        ((10, 5), (3, 4), (1, 3)),
+        ((10, 12), (2, 4), (1, 2)),
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype", [np.float32, np.float64, np.complex64, np.complex128],
+)
+def test_matmul(ashape, achunks, ablocks, bshape, bchunks, bblocks, dtype):
+    a = blosc2.linspace(0, 10, dtype=dtype, shape=ashape, chunks=achunks, blocks=ablocks)
+    b = blosc2.linspace(0, 10, dtype=dtype, shape=bshape, chunks=bchunks, blocks=bblocks)
+    blosc2_res = blosc2.matmul(a, b)
+
+    na = a[:]
+    nb = b[:]
+    np_res = np.matmul(na, nb)
+
+    np.testing.assert_allclose(blosc2_res, np_res, rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    ("ashape", "achunks", "ablocks"),
+    [
+        ((12, 11), (7, 5), (3, 1)),
+        ((0, 0), (0, 0), (0, 0)),
+        ((10,), (4,), (2,)),
+    ],
+)
+@pytest.mark.parametrize(
+    ("bshape", "bchunks", "bblocks"),
+    [
+        ((1, 5), (1, 4), (1, 3)),
+        ((4, 6), (2, 4), (1, 3)),
+        ((5,), (4,), (2,)),
+    ],
+)
+def test_matmul_shapes(ashape, achunks, ablocks, bshape, bchunks, bblocks):
+    a = blosc2.linspace(0, 10, shape=ashape, chunks=achunks, blocks=ablocks)
+    b = blosc2.linspace(0, 10, shape=bshape, chunks=bchunks, blocks=bblocks)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(a, b)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(b, a)
+
+
+@pytest.mark.parametrize("scalar", [
+    5,                      # int
+    5.3,                    # float
+    1 + 2j,                 # complex
+    np.int32(5),            # NumPy int32
+    np.int64(5),            # NumPy int64
+    np.float32(5.3),        # NumPy float32
+    np.float64(5.3),        # NumPy float64
+    np.complex64(1 + 2j),   # NumPy complex64
+    np.complex128(1 + 2j),  # NumPy complex128
+])
+def test_matmul_scalars(scalar):
+    vector = blosc2.asarray(np.array([1, 2, 3]))
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(scalar, vector)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(vector, scalar)
+
+    with pytest.raises(ValueError):
+        blosc2.matmul(scalar, scalar)