Efficient matrix, diagonal matrix multiplication for radial filters

Author: spors

micarray/modal/radial.py

@@ -193,10 +193,10 @@ def diagonal_mode_mat(bk):
     Returns
     -------
     Bk : (M, (N+1)**2, (N+1)**2) numpy.ndarray
-        Multidimensional array containing diagnonal matrices with input 
+        Multidimensional array containing diagnonal matrices with input
         vector on main diagonal.
     """
-    bk = _repeat_n_m(bk)
+    bk = repeat_n_m(bk)
     if len(bk.shape) == 1:
         bk = bk[np.newaxis, :]
     K, N = bk.shape
@@ -206,6 +206,24 @@ def diagonal_mode_mat(bk):
     return np.squeeze(Bk)
 
 
-def _repeat_n_m(v):
+def repeat_n_m(v):
+    """Repeat elements in a vector .
+
+    Returns a vector with the elements of the vector *v* repeated *n* times,
+    where *n* denotes the position of the element in *v*. The function can
+    be used to order the coefficients in the vector according to the order of
+    spherical harmonics. If *v* is a matrix, it is treated as a stack of
+    vectors residing in the last index and broadcast accordingly.
+
+    Parameters
+    ----------
+    v : (,N+1) numpy.ndarray
+        Input vector or stack of input vectors.
+
+    Returns
+    -------
+     : (,(N+1)**2) numpy.ndarray
+        Vector or stack of vectors containing repetated values.
+    """
     krlist = [np.tile(v, (2*i+1, 1)).T for i, v in enumerate(v.T.tolist())]
     return np.squeeze(np.concatenate(krlist, axis=-1))

micarray/util.py

@@ -60,3 +60,34 @@ def asarray_1d(a, **kwargs):
     elif result.ndim > 1:
         raise ValueError("array must be one-dimensional")
     return result
+
+
+def matdiagmul(A, b):
+    """Efficient multiplication of  matrix and diagonal matrix .
+
+    Returns the multiplication of a matrix *A* and a diagonal matrix. The
+    diagonal matrix is given by the vector *b* containing its elements on
+    the main diagonal.  If *b* is a matrix, it is treated as a stack of vectors
+    residing in the last index and broadcast accordingly.
+
+    Parameters
+    ----------
+    A : array_like
+        Input matrix.
+    b : array_like
+        Main diagonal elements or stack of main diagonal elements.
+
+    Returns
+    -------
+    array_like
+        Result of matrix multiplication.
+    """
+    if len(b.shape) == 1:
+        b = b[np.newaxis, :]
+    K, N = b.shape
+    M, N = A.shape
+
+    C = np.zeros([K, M, N], dtype=A.dtype)
+    for k in range(K):
+        C[k, :, :] = A * b[k, :]
+    return C