From 9808e0c3837e5fa8f2427749986e653cd9e052ba Mon Sep 17 00:00:00 2001
From: Edgar Solomonik <solomon2@illinois.edu>
Date: Sat, 20 Apr 2019 13:56:00 -0500
Subject: [PATCH] add norm() to function as numpy norm and fix up types yielded
 by operations on scalar-length arrays

---
 src_python/ctf/core.pyx | 68 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 62 insertions(+), 6 deletions(-)

diff --git a/src_python/ctf/core.pyx b/src_python/ctf/core.pyx
index 2b1fef7d..1ef81f6c 100644
--- a/src_python/ctf/core.pyx
+++ b/src_python/ctf/core.pyx
@@ -2552,13 +2552,13 @@ cdef class tensor:
     def __getitem__(self, key_init):
         [key, is_everything, is_single_val, is_contig, inds, corr_shape, one_shape] = _setgetitem_helper(self, key_init)
 
-        if is_everything:
-            return self.reshape(corr_shape)
-
         if is_single_val:
             vals = self.read(np.asarray([key]).reshape(1,self.ndim))
             return vals[0]
 
+        if is_everything:
+            return self.reshape(corr_shape)
+
         if is_contig:
             offs = [ind[0] for ind in inds]
             ends = [ind[1] for ind in inds]
@@ -5780,15 +5780,15 @@ def solve_tri(tensor L, tensor B, lower=True, from_left=True, transp_L=False):
 def vecnorm(A, ord=2):
     """
     vecnorm(A, ord=2)
-    Return norm of tensor A.
+    Return vector (elementwise) norm of tensor A.
 
     Parameters
     ----------
     A: tensor_like
-        Input tensor with 1 or 2 dimensions. If A is 1-D tensor, return a 2-D tensor with A on diagonal.
+        Input tensor with 1, 2 or more dimensions.
 
     ord: {int 1, 2, inf}, optional
-        Order of the norm.
+        Type the norm, 2=Frobenius.
 
     Returns
     -------
@@ -5816,6 +5816,62 @@ def vecnorm(A, ord=2):
     t_norm.stop()
     return nrm
 
+def norm(A, ord=2):
+    """
+    norm(A, ord='fro')
+    Return vector or matrix norm of tensor A.
+    If A a matrix, compute induced (1/2/infinity)-matrix norms or Frobenius norm, if A has one or more than three dimensions, treat as vector
+
+    Parameters
+    ----------
+    A: tensor_like
+        Input tensor with 1, 2, or more dimensions.
+
+    ord: {int 1, 2, inf}, optional
+        Order of the norm.
+
+    Returns
+    -------
+    output: tensor
+        Norm of tensor A.
+
+    Examples
+    --------
+    >>> import ctf
+    >>> import ctf.linalg as la
+    >>> a = ctf.astensor([3,4.])
+    >>> la.vecnorm(a)
+    5.0
+    """
+    t_norm = timer("pynorm")
+    t_norm.start()
+    if A.ndim == 2:
+        if ord == 'fro':
+            nrm = vecnorm(A)
+        elif ord == 2:
+            [U,S,VT] = svd(A,1)
+            nrm = S[0]
+        elif ord == 1:
+            nrm = max(sum(abs(A),axis=0))
+        elif ord == np.inf:
+            nrm = max(sum(abs(A),axis=1))
+        else:
+            raise ValueError('CTF PYTHON ERROR: CTF only supports 1/2/inf vector norms')
+    else:
+        if ord == 'fro':
+            nrm = A.norm2()
+        elif ord == 2:
+            nrm = A.norm2()
+        elif ord == 1:
+            nrm = A.norm1()
+        elif ord == np.inf:
+            nrm = A.norm_infty()
+        else:
+            raise ValueError('CTF PYTHON ERROR: CTF only supports 1/2/inf vector norms')
+    t_norm.stop()
+    return nrm
+
+
 def _match_tensor_types(first, other):
     if isinstance(first, tensor):
         tsr = first