API change for eigs (#593)

* fix call signature

* revert

* ???

* ????

* fix test

* fix argument order in matvec

* fix linter

* fix doc

* fix API

* linting

* linting

* tests fixed

* fix test

* fixing a bunch of things

* fix code

* check

* check added

* remove []

* remove []

Co-authored-by: Chase Roberts <chaseriley@google.com>

Author: mganahl

tensornetwork/backends/backend_test.py

@@ -316,7 +316,7 @@ def test_base_backend_eigs_not_implemented():
 def test_base_backend_eigs_lanczos_not_implemented():
   backend = BaseBackend()
   with pytest.raises(NotImplementedError):
-    backend.eigsh_lanczos(lambda x: x, [], np.ones((2)))
+    backend.eigsh_lanczos(lambda x: x, np.ones((2)))
 
 
 def test_base_backend_addition_not_implemented():
@@ -401,6 +401,8 @@ def test_base_backend_broadcast_left_multiplication_not_implemented():
   backend = BaseBackend()
   with pytest.raises(NotImplementedError):
     backend.broadcast_left_multiplication(np.ones((2, 2)), np.ones((2, 2)))
+
+
 def test_backend_instantiation(backend):
   backend1 = backend_factory.get_backend(backend)
   backend2 = backend_factory.get_backend(backend)

tensornetwork/backends/base_backend.py

@@ -334,24 +334,31 @@ def eigh(self, matrix: Tensor):
 
   def eigs(self,
            A: Callable,
+           args: Optional[List[Tensor]] = None,
            initial_state: Optional[Tensor] = None,
-           num_krylov_vecs: Optional[int] = 200,
-           numeig: Optional[int] = 1,
-           tol: Optional[float] = 1E-8,
-           which: Optional[Text] = 'LR',
-           maxiter: Optional[int] = None,
-           dtype: Optional[Type] = None) -> List[Tensor]:
-    """Arnoldi method for finding the lowest eigenvector-eigenvalue pairs of a
-    linear operator `A`. `A` can be either a linear operator type object or a
-    regular callable. If no `initial_state` is provided then `A` has to have an
-    attribute `shape` so that a suitable initial state can be randomly
-    generated.
+           shape: Optional[Tuple[int, ...]] = None,
+           dtype: Optional[Type[np.number]] = None,
+           num_krylov_vecs: int = 50,
+           numeig: int = 1,
+           tol: float = 1E-8,
+           which: Text = 'LR',
+           maxiter: Optional[int] = None) -> List[Tensor]:
+    """Arnoldi method for finding the lowest eigenvector-eigenvalue pairs 
+    of a linear operator `A`. `A` is a callable implementing the 
+    matrix-vector product. If no `initial_state` is provided then 
+    `shape` and `dtype` have to be passed so that a suitable initial 
+    state can be randomly  generated.
 
     Args:
       A: A (sparse) implementation of a linear operator
-      initial_state: An initial vector for the Lanczos algorithm. If `None`,
+      arsg: A list of arguments to `A`.  `A` will be called as
+        `res = A(initial_state, *args)`.
+      initial_state: An initial vector for the algorithm. If `None`,
         a random initial `Tensor` is created using the `numpy.random.randn`
         method.
+      shape: The shape of the input-dimension of `A`.
+      dtype: The dtype of the input `A`. If both no `initial_state` is provided,
+        a random initial state with shape `shape` and dtype `dtype` is created.
       num_krylov_vecs: The number of iterations (number of krylov vectors).
       numeig: The nummber of eigenvector-eigenvalue pairs to be computed.
         If `numeig > 1`, `reorthogonalize` has to be `True`.
@@ -374,7 +381,7 @@ def eigs(self,
 
   def eigsh_lanczos(self,
                     A: Callable,
-                    args: List,
+                    args: Optional[List[Tensor]] = None,
                     initial_state: Optional[Tensor] = None,
                     shape: Optional[Tuple[int, ...]] = None,
                     dtype: Optional[Type[np.number]] = None,

tensornetwork/backends/jax/jax_backend.py

@@ -230,7 +230,7 @@ def cmplx_random_uniform(complex_dtype, real_dtype):
   def eigsh_lanczos(
       self,
       A: Callable,
-      args: List[Tensor],
+      args: Optional[List[Tensor]] = None,
       initial_state: Optional[Tensor] = None,
       shape: Optional[Tuple] = None,
       dtype: Optional[Type[np.number]] = None,
@@ -307,6 +307,8 @@ def A(H,x):
        eigvals: A list of `numeig` lowest eigenvalues
        eigvecs: A list of `numeig` lowest eigenvectors
     """
+    if args is None:
+      args = []
     if num_krylov_vecs < numeig:
       raise ValueError('`num_krylov_vecs` >= `numeig` required!')
 

tensornetwork/backends/jax/jax_backend_test.py

@@ -414,30 +414,30 @@ def test_eigsh_lanczos_raises():
   backend = jax_backend.JaxBackend()
   with pytest.raises(
       ValueError, match='`num_krylov_vecs` >= `numeig` required!'):
-    backend.eigsh_lanczos(lambda x: x, [], numeig=10, num_krylov_vecs=9)
+    backend.eigsh_lanczos(lambda x: x, numeig=10, num_krylov_vecs=9)
   with pytest.raises(
       ValueError,
       match="Got numeig = 2 > 1 and `reorthogonalize = False`. "
       "Use `reorthogonalize=True` for `numeig > 1`"):
-    backend.eigsh_lanczos(lambda x: x, [], numeig=2, reorthogonalize=False)
+    backend.eigsh_lanczos(lambda x: x, numeig=2, reorthogonalize=False)
   with pytest.raises(
       ValueError,
       match="if no `initial_state` is passed, then `shape` and"
       "`dtype` have to be provided"):
-    backend.eigsh_lanczos(lambda x: x, [], shape=(10,), dtype=None)
+    backend.eigsh_lanczos(lambda x: x, shape=(10,), dtype=None)
   with pytest.raises(
       ValueError,
       match="if no `initial_state` is passed, then `shape` and"
       "`dtype` have to be provided"):
-    backend.eigsh_lanczos(lambda x: x, [], shape=None, dtype=np.float64)
+    backend.eigsh_lanczos(lambda x: x, shape=None, dtype=np.float64)
   with pytest.raises(
       ValueError,
       match="if no `initial_state` is passed, then `shape` and"
       "`dtype` have to be provided"):
-    backend.eigsh_lanczos(lambda x: x, [])
+    backend.eigsh_lanczos(lambda x: x)
   with pytest.raises(
       TypeError, match="Expected a `jax.array`. Got <class 'list'>"):
-    backend.eigsh_lanczos(lambda x: x, [], initial_state=[1, 2, 3])
+    backend.eigsh_lanczos(lambda x: x, initial_state=[1, 2, 3])
 
 
 @pytest.mark.parametrize("dtype", np_dtypes)

tensornetwork/backends/numpy/numpy_backend.py

@@ -178,13 +178,15 @@ def eigh(self, matrix: Tensor) -> Tuple[Tensor, Tensor]:
 
   def eigs(self,
            A: Callable,
+           args: Optional[List] = None,
            initial_state: Optional[Tensor] = None,
-           num_krylov_vecs: Optional[int] = 200,
-           numeig: Optional[int] = 6,
-           tol: Optional[float] = 1E-8,
-           which: Optional[Text] = 'LR',
-           maxiter: Optional[int] = None,
-           dtype: Optional[Type[np.number]] = None) -> Tuple[List, List]:
+           shape: Optional[Tuple[int, ...]] = None,
+           dtype: Optional[Type[np.number]] = None,
+           num_krylov_vecs: int = 50,
+           numeig: int = 6,
+           tol: float = 1E-8,
+           which: Text = 'LR',
+           maxiter: Optional[int] = None) -> Tuple[List, List]:
     """
     Arnoldi method for finding the lowest eigenvector-eigenvalue pairs
     of a linear operator `A`. `A` can be either a
@@ -196,9 +198,14 @@ def eigs(self,
 
     Args:
       A: A (sparse) implementation of a linear operator
-      initial_state: An initial vector for the Lanczos algorithm. If `None`,
+      args: A list of arguments to `A`.  `A` will be called as
+        `res = A(initial_state, *args)`.
+      initial_state: An initial vector for the algorithm. If `None`,
         a random initial `Tensor` is created using the `numpy.random.randn`
         method.
+      shape: The shape of the input-dimension of `A`.
+      dtype: The dtype of the input `A`. If both no `initial_state` is provided,
+        a random initial state with shape `shape` and dtype `dtype` is created.
       num_krylov_vecs: The number of iterations (number of krylov vectors).
       numeig: The nummber of eigenvector-eigenvalue pairs to be computed.
         If `numeig > 1`, `reorthogonalize` has to be `True`.
@@ -217,38 +224,37 @@ def eigs(self,
        `np.ndarray`: An array of `numeig` lowest eigenvalues
        `np.ndarray`: An array of `numeig` lowest eigenvectors
     """
+    if args is None:
+      args = []
     if which == 'SI':
       raise ValueError('which = SI is currently not supported.')
     if which == 'LI':
       raise ValueError('which = LI is currently not supported.')
 
-    if (initial_state is not None) and hasattr(A, 'shape'):
-      if initial_state.shape != A.shape[1]:
-        raise ValueError(
-            "A.shape[1]={} and initial_state.shape={} are incompatible.".format(
-                A.shape[1], initial_state.shape))
+    if numeig + 1 >= num_krylov_vecs:
+      raise ValueError('`num_krylov_vecs` > `numeig + 1` required!')
 
     if initial_state is None:
-      if not hasattr(A, 'shape'):
-        raise AttributeError("`A` has no  attribute `shape`. Cannot initialize "
-                             "lanczos. Please provide a valid `initial_state`")
-      if not hasattr(A, 'dtype'):
-        raise AttributeError(
-            "`A` has no  attribute `dtype`. Cannot initialize "
-            "lanczos. Please provide a valid `initial_state` with "
-            "a `dtype` attribute")
-
-      initial_state = self.randn(A.shape[1], A.dtype)
+      if (shape is None) or (dtype is None):
+        raise ValueError("if no `initial_state` is passed, then `shape` and"
+                         "`dtype` have to be provided")
+      initial_state = self.randn(shape, dtype)
 
     if not isinstance(initial_state, np.ndarray):
-      raise TypeError("Expected a `np.array`. Got {}".format(
+      raise TypeError("Expected a `np.ndarray`. Got {}".format(
           type(initial_state)))
+
+    shape = initial_state.shape
+
+    def matvec(vector):
+      return np.ravel(A(np.reshape(vector, shape), *args))
+
     #initial_state is an np.ndarray of rank 1, so we can
     #savely deduce the shape from it
     lop = sp.sparse.linalg.LinearOperator(
         dtype=initial_state.dtype,
-        shape=(initial_state.shape[0], initial_state.shape[0]),
-        matvec=A)
+        shape=(np.prod(initial_state.shape), np.prod(initial_state.shape)),
+        matvec=matvec)
     eta, U = sp.sparse.linalg.eigs(
         A=lop,
         k=numeig,
@@ -260,11 +266,11 @@ def eigs(self,
     if dtype:
       eta = eta.astype(dtype)
       U = U.astype(dtype)
-    return list(eta), [U[:, n] for n in range(numeig)]
+    return list(eta), [np.reshape(U[:, n], shape) for n in range(numeig)]
 
   def eigsh_lanczos(self,
                     A: Callable,
-                    args: List,
+                    args: Optional[List[Tensor]] = None,
                     initial_state: Optional[Tensor] = None,
                     shape: Optional[Tuple] = None,
                     dtype: Optional[Type[np.number]] = None,
@@ -309,6 +315,9 @@ def eigsh_lanczos(self,
        eigvals: A list of `numeig` lowest eigenvalues
        eigvecs: A list of `numeig` lowest eigenvectors
     """
+    if args is None:
+      args = []
+
     if num_krylov_vecs < numeig:
       raise ValueError('`num_krylov_vecs` >= `numeig` required!')