[python-dtcwt] 255/497: utils: add stacked matrix/{matrix, vector} products

[Ghislain Vaillant]

This is an automated email from the git hooks/post-receive script.

ghisvail-guest pushed a commit to branch debian/sid
in repository python-dtcwt.

commit 0884bcfd096666b75cd7cbdfaf16a3eac4a40d96
Author: Rich Wareham <rjw57 at cam.ac.uk>
Date:   Wed Jan 29 12:49:23 2014 +0000

    utils: add stacked matrix/{matrix,vector} products
    
    Add two functions which interpret multi-dimension arrays as either
    arrays of 2D matrices or arrays of vectors and find products between
    them.
    
    The functions make use of the numpy 'einsum' function which provides
    very fast products on large arrays.
---
 dtcwt/utils.py     | 45 +++++++++++++++++++++++++++
 tests/testutils.py | 89 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/dtcwt/utils.py b/dtcwt/utils.py
index c338eef..f18fd6b 100644
--- a/dtcwt/utils.py
+++ b/dtcwt/utils.py
@@ -124,3 +124,48 @@ def memoize(obj):
             cache[args] = obj(*args, **kwargs)
         return cache[args]
     return memoizer
+
+def stacked_2d_matrix_vector_prod(mats, vecs):
+    """
+    Interpret *mats* and *vecs* as arrays of 2D matrices and vectors. I.e.
+    *mats* has shape PxQxNxM and *vecs* has shape PxQxM. The result
+    is a PxQxN array equivalent to:
+
+    .. code::
+
+        result[i,j,:] = mats[i,j,:,:].dot(vecs[i,j,:])
+
+    for all valid row and column indices *i* and *j*.
+    """
+    return np.einsum('...ij,...j->...i', mats, vecs)
+
+def stacked_2d_vector_matrix_prod(vecs, mats):
+    """
+    Interpret *mats* and *vecs* as arrays of 2D matrices and vectors. I.e.
+    *mats* has shape PxQxNxM and *vecs* has shape PxQxN. The result
+    is a PxQxM array equivalent to:
+
+    .. code::
+
+        result[i,j,:] = mats[i,j,:,:].T.dot(vecs[i,j,:])
+
+    for all valid row and column indices *i* and *j*.
+    """
+    vecshape = np.array(vecs.shape + (1,))
+    vecshape[-1:-3:-1] = vecshape[-2:]
+    outshape = mats.shape[:-2] + (mats.shape[-1],)
+    return stacked_2d_matrix_matrix_prod(vecs.reshape(vecshape), mats).reshape(outshape)
+
+def stacked_2d_matrix_matrix_prod(mats1, mats2):
+    """
+    Interpret *mats1* and *mats2* as arrays of 2D matrices. I.e.
+    *mats1* has shape PxQxNxM and *mats2* has shape PxQxMxR. The result
+    is a PxQxNxR array equivalent to:
+
+    .. code::
+
+        result[i,j,:,:] = mats1[i,j,:,:].dot(mats2[i,j,:,:])
+
+    for all valid row and column indices *i* and *j*.
+    """
+    return np.einsum('...ij,...jk->...ik', mats1, mats2)
diff --git a/tests/testutils.py b/tests/testutils.py
index 8939b6a..46e49a5 100644
--- a/tests/testutils.py
+++ b/tests/testutils.py
@@ -1,6 +1,8 @@
+import itertools
 import numpy as np
+from six.moves import xrange
 
-from dtcwt.utils import drawedge, drawcirc, appropriate_complex_type_for
+from dtcwt.utils import *
 
 def test_complex_type_for_complex():
     assert np.issubsctype(appropriate_complex_type_for(np.zeros((2,3), np.complex64)), np.complex64)
@@ -19,3 +21,88 @@ def test_draw_circ():
 def test_draw_edge():
     e = drawedge(20, (20,20), 4, 50)
     assert e.shape == (50, 50)
+
+def test_stacked_2d_matrix_vector_product():
+    nr, nc = (30, 20)
+    mr, mc = (3, 4)
+
+    # Generate a random stack of 2d matrices and 2d vectors
+    matrices = np.random.rand(nr, nc, mr, mc)
+    vectors = np.random.rand(nr, nc, mc)
+
+    # Compute product to test
+    out = stacked_2d_matrix_vector_prod(matrices, vectors)
+
+    # Check output shape is what we expect
+    assert out.shape == (nr, nc, mr)
+
+    # Check each product
+    for r, c in itertools.product(xrange(nr), xrange(nc)):
+        m = matrices[r, c, :, :]
+        v = vectors[r, c, :]
+        o = out[r, c, :]
+
+        assert m.shape == (mr, mc)
+        assert v.shape == (mc,)
+        assert o.shape == (mr,)
+
+        gold = m.dot(v)
+        max_delta = np.abs(gold - o).max()
+        assert max_delta < 1e-8
+
+def test_stacked_2d_vector_matrix_product():
+    nr, nc = (30, 20)
+    mr, mc = (3, 4)
+
+    # Generate a random stack of 2d matrices and 2d vectors
+    matrices = np.random.rand(nr, nc, mr, mc)
+    vectors = np.random.rand(nr, nc, mr)
+
+    # Compute product to test
+    out = stacked_2d_vector_matrix_prod(vectors, matrices)
+
+    # Check output shape is what we expect
+    assert out.shape == (nr, nc, mc)
+
+    # Check each product
+    for r, c in itertools.product(xrange(nr), xrange(nc)):
+        m = matrices[r, c, :, :]
+        v = vectors[r, c, :]
+        o = out[r, c, :]
+
+        assert m.shape == (mr, mc)
+        assert v.shape == (mr,)
+        assert o.shape == (mc,)
+
+        gold = v.dot(m)
+        max_delta = np.abs(gold - o).max()
+        assert max_delta < 1e-8
+
+def test_stacked_2d_matrix_matrix_product():
+    nr, nc = (30, 20)
+    mr1, k, mc2 = (3, 4, 5)
+
+    # Generate a random stack of 2d matrices and 2d vectors
+    matrices1 = np.random.rand(nr, nc, mr1, k)
+    matrices2 = np.random.rand(nr, nc, k, mc2)
+
+    # Compute product to test
+    out = stacked_2d_matrix_matrix_prod(matrices1, matrices2)
+
+    # Check output shape is what we expect
+    assert out.shape == (nr, nc, mr1, mc2)
+
+    # Check each product
+    for r, c in itertools.product(xrange(nr), xrange(nc)):
+        m1 = matrices1[r, c, :, :]
+        m2 = matrices2[r, c, :, :]
+        o = out[r, c, :, :]
+
+        assert m1.shape == (mr1, k)
+        assert m2.shape == (k, mc2)
+        assert o.shape == (mr1, mc2)
+
+        gold = m1.dot(m2)
+        max_delta = np.abs(gold - o).max()
+        assert max_delta < 1e-8
+

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-science/packages/python-dtcwt.git