[pypy-commit] pypy default: add numpypy.vdot

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r61959:3778a9fd5fe9
Date: 2013-03-03 13:57 -0500
http://bitbucket.org/pypy/pypy/changeset/3778a9fd5fe9/

Log:	add numpypy.vdot

diff --git a/lib_pypy/numpypy/core/numeric.py b/lib_pypy/numpypy/core/numeric.py
--- a/lib_pypy/numpypy/core/numeric.py
+++ b/lib_pypy/numpypy/core/numeric.py
@@ -2,7 +2,7 @@
            'newaxis', 'ufunc',
            'asarray', 'asanyarray', 'base_repr',
            'array_repr', 'array_str', 'set_string_function',
-           'array_equal', 'outer', 'identity', 'little_endian',
+           'array_equal', 'outer', 'vdot', 'identity', 'little_endian',
            'Inf', 'inf', 'infty', 'Infinity', 'nan', 'NaN', 'False_', 'True_',
           ]
 
@@ -523,6 +523,60 @@
     b = asarray(b)
     return a.ravel()[:,newaxis]*b.ravel()[newaxis,:]
 
+def vdot(a, b):
+    """
+    Return the dot product of two vectors.
+
+    The vdot(`a`, `b`) function handles complex numbers differently than
+    dot(`a`, `b`).  If the first argument is complex the complex conjugate
+    of the first argument is used for the calculation of the dot product.
+
+    Note that `vdot` handles multidimensional arrays differently than `dot`:
+    it does *not* perform a matrix product, but flattens input arguments
+    to 1-D vectors first. Consequently, it should only be used for vectors.
+
+    Parameters
+    ----------
+    a : array_like
+        If `a` is complex the complex conjugate is taken before calculation
+        of the dot product.
+    b : array_like
+        Second argument to the dot product.
+
+    Returns
+    -------
+    output : ndarray
+        Dot product of `a` and `b`.  Can be an int, float, or
+        complex depending on the types of `a` and `b`.
+
+    See Also
+    --------
+    dot : Return the dot product without using the complex conjugate of the
+          first argument.
+
+    Examples
+    --------
+    >>> a = np.array([1+2j,3+4j])
+    >>> b = np.array([5+6j,7+8j])
+    >>> np.vdot(a, b)
+    (70-8j)
+    >>> np.vdot(b, a)
+    (70+8j)
+
+    Note that higher-dimensional arrays are flattened!
+
+    >>> a = np.array([[1, 4], [5, 6]])
+    >>> b = np.array([[4, 1], [2, 2]])
+    >>> np.vdot(a, b)
+    30
+    >>> np.vdot(b, a)
+    30
+    >>> 1*4 + 4*1 + 5*2 + 6*2
+    30
+
+    """
+    return dot(asarray(a).ravel().conj(), asarray(b).ravel())
+
 def identity(n, dtype=None):
     """
     Return the identity array.
diff --git a/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py b/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py
--- a/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py
+++ b/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py
@@ -189,6 +189,22 @@
                           [12, 15, 18]])
         assert (res == expected).all()
 
+    def test_vdot(self):
+        import numpypy as np
+        a = np.array([1+2j,3+4j])
+        b = np.array([5+6j,7+8j])
+        c = np.vdot(a, b)
+        assert c == (70-8j)
+        c = np.vdot(b, a)
+        assert c == (70+8j)
+
+        a = np.array([[1, 4], [5, 6]])
+        b = np.array([[4, 1], [2, 2]])
+        c = np.vdot(a, b)
+        assert c == 30
+        c = np.vdot(b, a)
+        assert c == 30
+
     def test_identity(self):
         from numpypy import array, int32, float64, dtype, identity
         a = identity(0)