Merge pull request #3387 from WarrenWeckesser/norm-axis

ENH: linalg: Add an `axis` argument to linalg.norm

2 files changed, 209 insertions, 51 deletions

diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py
index ae0da3685..f11f905f7 100644
--- a/numpy/linalg/linalg.py
+++ b/numpy/linalg/linalg.py
@@ -20,10 +20,10 @@ import warnings
 
 from numpy.core import array, asarray, zeros, empty, transpose, \
         intc, single, double, csingle, cdouble, inexact, complexfloating, \
-        newaxis, ravel, all, Inf, dot, add, multiply, identity, sqrt, \
-        maximum, flatnonzero, diagonal, arange, fastCopyAndTranspose, sum, \
-        isfinite, size, finfo, absolute, log, exp, errstate, geterrobj
-from numpy.lib import triu
+        newaxis, ravel, all, Inf, dot, add, multiply, sqrt, maximum, \
+        fastCopyAndTranspose, sum, isfinite, size, finfo, errstate, \
+        geterrobj, float128, rollaxis, amin, amax
+from numpy.lib import triu, asfarray
 from numpy.linalg import lapack_lite, _umath_linalg
 from numpy.matrixlib.defmatrix import matrix_power
 from numpy.compat import asbytes
@@ -1865,7 +1865,38 @@ def lstsq(a, b, rcond=-1):
     st = s[:min(n, m)].copy().astype(result_real_t)
     return wrap(x), wrap(resids), results['rank'], st
 
-def norm(x, ord=None):
+
+def _multi_svd_norm(x, row_axis, col_axis, op):
+    """Compute the exteme singular values of the 2-D matrices in `x`.
+
+    This is a private utility function used by numpy.linalg.norm().
+
+    Parameters
+    ----------
+    x : ndarray
+    row_axis, col_axis : int
+        The axes of `x` that hold the 2-D matrices.
+    op : callable
+        This should be either numpy.amin or numpy.amax.
+
+    Returns
+    -------
+    result : float or ndarray
+        If `x` is 2-D, the return values is a float.
+        Otherwise, it is an array with ``x.ndim - 2`` dimensions.
+        The return values are either the minimum or maximum of the
+        singular values of the matrices, depending on whether `op`
+        is `numpy.amin` or `numpy.amax`.
+
+    """
+    if row_axis > col_axis:
+        row_axis -= 1
+    y = rollaxis(rollaxis(x, col_axis, x.ndim), row_axis, -1)
+    result = op(svd(y, compute_uv=0), axis=-1)
+    return result
+
+
+def norm(x, ord=None, axis=None):
     """
     Matrix or vector norm.
 
@@ -1875,16 +1906,22 @@ def norm(x, ord=None):
 
     Parameters
     ----------
-    x : {(M,), (M, N)} array_like
-        Input array.
+    x : array_like
+        Input array.  If `axis` is None, `x` must be 1-D or 2-D.
     ord : {non-zero int, inf, -inf, 'fro'}, optional
         Order of the norm (see table under ``Notes``). inf means numpy's
         `inf` object.
+    axis : {int, 2-tuple of ints, None}, optional
+        If `axis` is an integer, it specifies the axis of `x` along which to
+        compute the vector norms.  If `axis` is a 2-tuple, it specifies the
+        axes that hold 2-D matrices, and the matrix norms of these matrices
+        are computed.  If `axis` is None then either a vector norm (when `x`
+        is 1-D) or a matrix norm (when `x` is 2-D) is returned.
 
     Returns
     -------
-    n : float
-        Norm of the matrix or vector.
+    n : float or ndarray
+        Norm of the matrix or vector(s).
 
     Notes
     -----
@@ -1967,44 +2004,96 @@ def norm(x, ord=None):
     >>> LA.norm(a, -3)
     nan
 
+    Using the `axis` argument to compute vector norms:
+
+    >>> c = np.array([[ 1, 2, 3],
+    ...               [-1, 1, 4]])
+    >>> LA.norm(c, axis=0)
+    array([ 1.41421356,  2.23606798,  5.        ])
+    >>> LA.norm(c, axis=1)
+    array([ 3.74165739,  4.24264069])
+    >>> LA.norm(c, ord=1, axis=1)
+    array([6, 6])
+
+    Using the `axis` argument to compute matrix norms:
+
+    >>> m = np.arange(8).reshape(2,2,2)
+    >>> LA.norm(m, axis=(1,2))
+    array([  3.74165739,  11.22497216])
+    >>> LA.norm(m[0, :, :]), LA.norm(m[1, :, :])
+    (3.7416573867739413, 11.224972160321824)
+
     """
     x = asarray(x)
-    if ord is None: # check the default case first and handle it immediately
+
+    # Check the default case first and handle it immediately.
+    if ord is None and axis is None:
+        s = (x.conj() * x).real
         return sqrt(add.reduce((x.conj() * x).ravel().real))
 
+    # Normalize the `axis` argument to a tuple.
+    if axis is None:
+        axis = tuple(range(x.ndim))
+    elif not isinstance(axis, tuple):
+        axis = (axis,)
+
     nd = x.ndim
-    if nd == 1:
+    if len(axis) == 1:
         if ord == Inf:
-            return abs(x).max()
+            return abs(x).max(axis=axis)
         elif ord == -Inf:
-            return abs(x).min()
+            return abs(x).min(axis=axis)
         elif ord == 0:
-            return (x != 0).sum() # Zero norm
+            # Zero norm
+            return (x != 0).sum(axis=axis)
         elif ord == 1:
-            return abs(x).sum() # special case for speedup
-        elif ord == 2:
-            return sqrt(((x.conj()*x).real).sum()) # special case for speedup
+            # special case for speedup
+            return add.reduce(abs(x), axis=axis)
+        elif ord is None or ord == 2:
+            # special case for speedup
+            s = (x.conj() * x).real
+            return sqrt(add.reduce(s, axis=axis))
         else:
             try:
                 ord + 1
             except TypeError:
                 raise ValueError("Invalid norm order for vectors.")
-            return ((abs(x)**ord).sum())**(1.0/ord)
-    elif nd == 2:
+            if x.dtype != float128:
+                # Convert to a float type, so integer arrays give
+                # float results.  Don't apply asfarray to float128 arrays,
+                # because it will downcast to float64.
+                absx = asfarray(abs(x))
+            return add.reduce(absx**ord, axis=axis)**(1.0/ord)
+    elif len(axis) == 2:
+        row_axis, col_axis = axis
+        if not (-x.ndim <= row_axis < x.ndim and
+                -x.ndim <= col_axis < x.ndim):
+            raise ValueError('Invalid axis %r for an array with shape %r' %
+                             (axis, x.shape))
+        if row_axis % x.ndim == col_axis % x.ndim:
+            raise ValueError('Duplicate axes given.')
         if ord == 2:
-            return svd(x, compute_uv=0).max()
+            return _multi_svd_norm(x, row_axis, col_axis, amax)
         elif ord == -2:
-            return svd(x, compute_uv=0).min()
+            return _multi_svd_norm(x, row_axis, col_axis, amin)
         elif ord == 1:
-            return abs(x).sum(axis=0).max()
+            if col_axis > row_axis:
+                col_axis -= 1
+            return add.reduce(abs(x), axis=row_axis).max(axis=col_axis)
         elif ord == Inf:
-            return abs(x).sum(axis=1).max()
+            if row_axis > col_axis:
+                row_axis -= 1
+            return add.reduce(abs(x), axis=col_axis).max(axis=row_axis)
         elif ord == -1:
-            return abs(x).sum(axis=0).min()
+            if col_axis > row_axis:
+                col_axis -= 1
+            return add.reduce(abs(x), axis=row_axis).min(axis=col_axis)
         elif ord == -Inf:
-            return abs(x).sum(axis=1).min()
-        elif ord in ['fro','f']:
-            return sqrt(add.reduce((x.conj() * x).real.ravel()))
+            if row_axis > col_axis:
+                row_axis -= 1
+            return add.reduce(abs(x), axis=col_axis).min(axis=row_axis)
+        elif ord in [None, 'fro', 'f']:
+            return sqrt(add.reduce((x.conj() * x).real, axis=axis))
         else:
             raise ValueError("Invalid norm order for matrices.")
     else:
diff --git a/numpy/linalg/tests/test_linalg.py b/numpy/linalg/tests/test_linalg.py
index 2dc55ab5e..881311c94 100644
--- a/numpy/linalg/tests/test_linalg.py
+++ b/numpy/linalg/tests/test_linalg.py
@@ -515,30 +515,37 @@ class TestEigh(HermitianTestCase, HermitianGeneralizedTestCase, TestCase):
 
 
 class _TestNorm(TestCase):
+
     dt = None
     dec = None
+
     def test_empty(self):
         assert_equal(norm([]), 0.0)
         assert_equal(norm(array([], dtype=self.dt)), 0.0)
         assert_equal(norm(atleast_2d(array([], dtype=self.dt))), 0.0)
 
     def test_vector(self):
-        a = [1.0,2.0,3.0,4.0]
-        b = [-1.0,-2.0,-3.0,-4.0]
-        c = [-1.0, 2.0,-3.0, 4.0]
+        a = [1, 2, 3, 4]
+        b = [-1, -2, -3, -4]
+        c = [-1, 2, -3, 4]
 
         def _test(v):
-            np.testing.assert_almost_equal(norm(v), 30**0.5, decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,inf), 4.0, decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,-inf), 1.0, decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,1), 10.0, decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,-1), 12.0/25,
-                    decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,2), 30**0.5,
-                    decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,-2), ((205./144)**-0.5),
-                    decimal=self.dec)
-            np.testing.assert_almost_equal(norm(v,0), 4, decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v), 30**0.5,
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, inf), 4.0,
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, -inf), 1.0,
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, 1), 10.0,
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, -1), 12.0/25,
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, 2), 30**0.5,
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, -2), ((205./144)**-0.5),
+                                           decimal=self.dec)
+            np.testing.assert_almost_equal(norm(v, 0), 4,
+                                           decimal=self.dec)
 
         for v in (a, b, c,):
             _test(v)
@@ -548,25 +555,82 @@ class _TestNorm(TestCase):
             _test(v)
 
     def test_matrix(self):
-        A = matrix([[1.,3.],[5.,7.]], dtype=self.dt)
-        A = matrix([[1.,3.],[5.,7.]], dtype=self.dt)
+        A = matrix([[1, 3], [5, 7]], dtype=self.dt)
         assert_almost_equal(norm(A), 84**0.5)
-        assert_almost_equal(norm(A,'fro'), 84**0.5)
-        assert_almost_equal(norm(A,inf), 12.0)
-        assert_almost_equal(norm(A,-inf), 4.0)
-        assert_almost_equal(norm(A,1), 10.0)
-        assert_almost_equal(norm(A,-1), 6.0)
-        assert_almost_equal(norm(A,2), 9.1231056256176615)
-        assert_almost_equal(norm(A,-2), 0.87689437438234041)
+        assert_almost_equal(norm(A, 'fro'), 84**0.5)
+        assert_almost_equal(norm(A, inf), 12.0)
+        assert_almost_equal(norm(A, -inf), 4.0)
+        assert_almost_equal(norm(A, 1), 10.0)
+        assert_almost_equal(norm(A, -1), 6.0)
+        assert_almost_equal(norm(A, 2), 9.1231056256176615)
+        assert_almost_equal(norm(A, -2), 0.87689437438234041)
 
         self.assertRaises(ValueError, norm, A, 'nofro')
         self.assertRaises(ValueError, norm, A, -3)
         self.assertRaises(ValueError, norm, A, 0)
 
+    def test_axis(self):
+        # Vector norms.
+        # Compare the use of `axis` with computing the norm of each row
+        # or column separately.
+        A = array([[1, 2, 3], [4, 5, 6]], dtype=self.dt)
+        for order in [None, -1, 0, 1, 2, 3, np.Inf, -np.Inf]:
+            expected0 = [norm(A[:,k], ord=order) for k in range(A.shape[1])]
+            assert_almost_equal(norm(A, ord=order, axis=0), expected0)
+            expected1 = [norm(A[k,:], ord=order) for k in range(A.shape[0])]
+            assert_almost_equal(norm(A, ord=order, axis=1), expected1)
+
+        # Matrix norms.
+        B = np.arange(1, 25, dtype=self.dt).reshape(2, 3, 4)
+
+        for order in [None, -2, 2, -1, 1, np.Inf, -np.Inf, 'fro']:
+            assert_almost_equal(norm(A, ord=order), norm(A, ord=order,
+                                                         axis=(0, 1)))
+
+            n = norm(B, ord=order, axis=(1, 2))
+            expected = [norm(B[k], ord=order) for k in range(B.shape[0])]
+            assert_almost_equal(n, expected)
+
+            n = norm(B, ord=order, axis=(2, 1))
+            expected = [norm(B[k].T, ord=order) for k in range(B.shape[0])]
+            assert_almost_equal(n, expected)
+
+            n = norm(B, ord=order, axis=(0, 2))
+            expected = [norm(B[:,k,:], ord=order) for k in range(B.shape[1])]
+            assert_almost_equal(n, expected)
+
+            n = norm(B, ord=order, axis=(0, 1))
+            expected = [norm(B[:,:,k], ord=order) for k in range(B.shape[2])]
+            assert_almost_equal(n, expected)
+
+    def test_bad_args(self):
+        # Check that bad arguments raise the appropriate exceptions.
+
+        A = array([[1, 2, 3], [4, 5, 6]], dtype=self.dt)
+        B = np.arange(1, 25, dtype=self.dt).reshape(2, 3, 4)
+
+        # Using `axis=<integer>` or passing in a 1-D array implies vector
+        # norms are being computed, so also using `ord='fro'` raises a
+        # ValueError.
+        self.assertRaises(ValueError, norm, A, 'fro', 0)
+        self.assertRaises(ValueError, norm, [3, 4], 'fro', None)
+
+        # Similarly, norm should raise an exception when ord is any finite
+        # number other than 1, 2, -1 or -2 when computing matrix norms.
+        for order in [0, 3]:
+            self.assertRaises(ValueError, norm, A, order, None)
+            self.assertRaises(ValueError, norm, A, order, (0, 1))
+            self.assertRaises(ValueError, norm, B, order, (1, 2))
+
+        # Invalid axis
+        self.assertRaises(ValueError, norm, B, None, 3)
+        self.assertRaises(ValueError, norm, B, None, (2, 3))
+        self.assertRaises(ValueError, norm, B, None, (0, 1, 2))
+
 
 class TestNormDouble(_TestNorm):
     dt = np.double
-    dec= 12
+    dec = 12
 
 
 class TestNormSingle(_TestNorm):
@@ -574,6 +638,11 @@ class TestNormSingle(_TestNorm):
     dec = 6
 
 
+class TestNormInt64(_TestNorm):
+    dt = np.int64
+    dec = 12
+
+
 class TestMatrixRank(object):
     def test_matrix_rank(self):
         # Full rank matrix