Eigsh (#592)

* fix call signature

* revert

* ???

* ????

* fix test

* fix argument order in matvec

* fix linter

* fix doc

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

Author: mganahl

tensornetwork/backends/base_backend.py

@@ -389,10 +389,10 @@ def eigsh_lanczos(self,
     of `A`.
     Args:
       A: A (sparse) implementation of a linear operator.
-         Call signature of `A` is `res = A(*args, vector)`, where `vector`
+         Call signature of `A` is `res = A(vector, *args)`, where `vector`
          can be an arbitrary `Tensor`, and `res.shape` has to be `vector.shape`.
       arsg: A list of arguments to `A`.  `A` will be called as
-        `res = A(*args, initial_state)`.
+        `res = A(initial_state, *args)`.
       initial_state: An initial vector for the Lanczos algorithm. If `None`,
         a random initial `Tensor` is created using the `backend.randn` method
       shape: The shape of the input-dimension of `A`.

tensornetwork/backends/jax/jax_backend.py

@@ -277,10 +277,10 @@ def A(H,x):
     
     Args:
       A: A (sparse) implementation of a linear operator.
-         Call signature of `A` is `res = A(*args, vector)`, where `vector`
+         Call signature of `A` is `res = A(vector, *args)`, where `vector`
          can be an arbitrary `Tensor`, and `res.shape` has to be `vector.shape`.
       arsg: A list of arguments to `A`.  `A` will be called as
-        `res = A(*args, initial_state)`.
+        `res = A(initial_state, *args)`.
       initial_state: An initial vector for the Lanczos algorithm. If `None`,
         a random initial `Tensor` is created using the `backend.randn` method
       shape: The shape of the input-dimension of `A`.

tensornetwork/backends/jax/jax_backend_test.py

@@ -308,10 +308,10 @@ def test_eigsh_valid_init_operator_with_shape(dtype):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def matvec(H, x):
+  def mv(x, H):
     return jax.numpy.dot(H, x)
 
-  eta1, U1 = backend.eigsh_lanczos(matvec, [H], init)
+  eta1, U1 = backend.eigsh_lanczos(mv, [H], init)
   eta2, U2 = np.linalg.eigh(H)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
@@ -326,10 +326,10 @@ def test_eigsh_small_number_krylov_vectors():
   init = np.array([1, 1], dtype=np.float64)
   H = np.array([[1, 2], [3, 4]], dtype=np.float64)
 
-  def matvec(H, x):
+  def mv(x, H):
     return jax.numpy.dot(H, x)
 
-  eta1, _ = backend.eigsh_lanczos(matvec, [H], init, num_krylov_vecs=1)
+  eta1, _ = backend.eigsh_lanczos(mv, [H], init, num_krylov_vecs=1)
   np.testing.assert_allclose(eta1[0], 5)
 
 
@@ -342,7 +342,7 @@ def test_eigsh_lanczos_1(dtype):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(H, x):
+  def mv(x, H):
     return jax.numpy.dot(H, x)
 
   eta1, U1 = backend.eigsh_lanczos(mv, [H], init)
@@ -363,7 +363,7 @@ def test_eigsh_lanczos_2(dtype):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(H, x):
+  def mv(x, H):
     return jax.numpy.dot(H, x)
 
   eta1, U1 = backend.eigsh_lanczos(mv, [H], shape=(D,), dtype=dtype)
@@ -385,11 +385,11 @@ def test_eigsh_lanczos_reorthogonalize(dtype, numeig):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def matvec(H, x):
+  def mv(x, H):
     return jax.numpy.dot(H, x)
 
   eta1, U1 = backend.eigsh_lanczos(
-      matvec, [H],
+      mv, [H],
       shape=(D,),
       dtype=dtype,
       numeig=numeig,

tensornetwork/backends/jax/jitted_functions.py

@@ -19,7 +19,7 @@ def _generate_jitted_eigsh_lanczos(jax):
   ```
   `matvec`: A callable implementing the matrix-vector product of a 
   linear operator. `arguments`: Arguments to `matvec` additional to 
-  an input vector. `matvec` will be called as `matvec(*args, init)`.
+  an input vector. `matvec` will be called as `matvec(init, *args)`.
   `init`: An initial input state to `matvec`.
   `ncv`: Number of krylov iterations (i.e. dimension of the Krylov space).
   `neig`: Number of eigenvalue-eigenvector pairs to be computed.
@@ -43,7 +43,7 @@ def jax_lanczos(matvec, arguments, init, ncv, neig, landelta, reortho):
       matvec: A callable implementing the matrix-vector product of a 
         linear operator.
       arguments: Arguments to `matvec` additional to an input vector. 
-        `matvec` will be called as `matvec(*args, init)`.
+        `matvec` will be called as `matvec(init, *args)`.
       init: An initial input state to `matvec`.
       ncv: Number of krylov iterations (i.e. dimension of the Krylov space).
       neig: Number of eigenvalue-eigenvector pairs to be computed.
@@ -75,7 +75,7 @@ def body_lanczos(vals):
           lambda x: jax.lax.fori_loop(0, i, body_modified_gram_schmidt,
                                       [normalized_vector, krylov_vectors]),
           False, lambda x: [normalized_vector, krylov_vectors])
-      Av = matvec(*args, normalized_vector)
+      Av = matvec(normalized_vector, *args)
 
       diag_element = jax.numpy.dot(
           jax.numpy.conj(jax.numpy.ravel(normalized_vector)),

tensornetwork/backends/numpy/numpy_backend.py

@@ -279,10 +279,10 @@ def eigsh_lanczos(self,
     of a linear operator `A`.
     Args:
       A: A (sparse) implementation of a linear operator.
-         Call signature of `A` is `res = A(*args, vector)`, where `vector`
+         Call signature of `A` is `res = A(vector, *args)`, where `vector`
          can be an arbitrary `Tensor`, and `res.shape` has to be `vector.shape`.
       arsg: A list of arguments to `A`.  `A` will be called as
-        `res = A(*args, initial_state)`.
+        `res = A(initial_state, *args)`.
       initial_state: An initial vector for the Lanczos algorithm. If `None`,
         a random initial `Tensor` is created using the `backend.randn` method
       shape: The shape of the input-dimension of `A`.
@@ -346,7 +346,7 @@ def eigsh_lanczos(self,
         for v in krylov_vecs:
           vector_n -= np.dot(np.ravel(np.conj(v)), np.ravel(vector_n)) * v
       krylov_vecs.append(vector_n)
-      A_vector_n = A(*args, vector_n)
+      A_vector_n = A(vector_n, *args)
       diag_elements.append(
           np.dot(np.ravel(np.conj(vector_n)), np.ravel(A_vector_n)))
 

tensornetwork/backends/numpy/numpy_backend_test.py

@@ -287,10 +287,10 @@ def test_eigsh_valid_init_operator_with_shape(dtype):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return np.dot(H, x)
+  def mv(x, mat):
+    return np.dot(mat, x)
 
-  eta1, U1 = backend.eigsh_lanczos(mv, [], init)
+  eta1, U1 = backend.eigsh_lanczos(mv, [H], init)
   eta2, U2 = np.linalg.eigh(H)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
@@ -305,10 +305,10 @@ def test_eigsh_small_number_krylov_vectors():
   init = np.array([1, 1], dtype=np.float64)
   H = np.array([[1, 2], [3, 4]], dtype=np.float64)
 
-  def mv(x):
-    return np.dot(H, x)
+  def mv(x, mat):
+    return np.dot(mat, x)
 
-  eta1, _ = backend.eigsh_lanczos(mv, [], init, num_krylov_vecs=1)
+  eta1, _ = backend.eigsh_lanczos(mv, [H], init, num_krylov_vecs=1)
   np.testing.assert_allclose(eta1[0], 5)
 
 
@@ -321,10 +321,10 @@ def test_eigsh_lanczos_1(dtype):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return np.dot(H, x)
+  def mv(x, mat):
+    return np.dot(mat, x)
 
-  eta1, U1 = backend.eigsh_lanczos(mv, [], init)
+  eta1, U1 = backend.eigsh_lanczos(mv, [H], init)
   eta2, U2 = np.linalg.eigh(H)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
@@ -342,10 +342,10 @@ def test_eigsh_lanczos_2(dtype):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return np.dot(H, x)
+  def mv(x, mat):
+    return np.dot(mat, x)
 
-  eta1, U1 = backend.eigsh_lanczos(mv, [], shape=(D,), dtype=dtype)
+  eta1, U1 = backend.eigsh_lanczos(mv, [H], shape=(D,), dtype=dtype)
   eta2, U2 = np.linalg.eigh(H)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
@@ -364,11 +364,11 @@ def test_eigsh_lanczos_reorthogonalize(dtype, numeig):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return np.dot(H, x)
+  def mv(x, mat):
+    return np.dot(mat, x)
 
   eta1, U1 = backend.eigsh_lanczos(
-      mv, [],
+      mv, [H],
       shape=(D,),
       dtype=dtype,
       numeig=numeig,

tensornetwork/backends/pytorch/pytorch_backend.py

@@ -187,10 +187,10 @@ def eigsh_lanczos(self,
     of a `LinearOperator` `A`.
     Args:
       A: A (sparse) implementation of a linear operator.
-         Call signature of `A` is `res = A(*args, vector)`, where `vector`
+         Call signature of `A` is `res = A(vector, *args)`, where `vector`
          can be an arbitrary `Tensor`, and `res.shape` has to be `vector.shape`.
       arsg: A list of arguments to `A`.  `A` will be called as
-        `res = A(*args, initial_state)`.
+        `res = A(initial_state, *args)`.
       initial_state: An initial vector for the Lanczos algorithm. If `None`,
         a random initial `Tensor` is created using the `torch.randn` method
       shape: The shape of the input-dimension of `A`.
@@ -257,7 +257,7 @@ def eigsh_lanczos(self,
           vector_n -= (v.view(-1).dot(vector_n.view(-1))) * torchlib.reshape(
               v, vector_n.shape)
       krylov_vecs.append(vector_n)
-      A_vector_n = A(*args, vector_n)
+      A_vector_n = A(vector_n, *args)
       diag_elements.append(vector_n.view(-1).dot(A_vector_n.view(-1)))
 
       if ((it > 0) and (it % ndiag) == 0) and (len(diag_elements) >= numeig):

tensornetwork/backends/pytorch/pytorch_backend_test.py

@@ -259,10 +259,10 @@ def test_eigsh_lanczos_1():
   tmp = backend.randn((D, D), dtype=dtype)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return H.mv(x)
+  def mv(x, mat):
+    return mat.mv(x)
 
-  eta1, U1 = backend.eigsh_lanczos(mv, [], init, num_krylov_vecs=D)
+  eta1, U1 = backend.eigsh_lanczos(mv, [H], init, num_krylov_vecs=D)
   eta2, U2 = H.symeig(eigenvectors=True)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
@@ -277,10 +277,10 @@ def test_eigsh_small_number_krylov_vectors():
   init = backend.convert_to_tensor(np.array([1, 1], dtype=np.float64))
   H = backend.convert_to_tensor(np.array([[1, 2], [3, 4]], dtype=np.float64))
 
-  def mv(x):
-    return H.mv(x)
+  def mv(x, mat):
+    return mat.mv(x)
 
-  eta1, _ = backend.eigsh_lanczos(mv, [], init, num_krylov_vecs=1)
+  eta1, _ = backend.eigsh_lanczos(mv, [H], init, num_krylov_vecs=1)
   np.testing.assert_allclose(eta1[0], 5)
 
 
@@ -293,11 +293,11 @@ def test_eigsh_lanczos_reorthogonalize(numeig):
   tmp = backend.randn((D, D), dtype=dtype, seed=10)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return H.mv(x)
+  def mv(x, mat):
+    return mat.mv(x)
 
   eta1, U1 = backend.eigsh_lanczos(
-      mv, [],
+      mv, [H],
       shape=(D,),
       dtype=dtype,
       numeig=numeig,
@@ -325,11 +325,11 @@ def test_eigsh_lanczos_2():
   tmp = backend.randn((D, D), dtype=dtype)
   H = tmp + backend.transpose(backend.conj(tmp), (1, 0))
 
-  def mv(x):
-    return H.mv(x)
+  def mv(x, mat):
+    return mat.mv(x)
 
   eta1, U1 = backend.eigsh_lanczos(
-      mv, [],
+      mv, [H],
       shape=(D,),
       dtype=dtype,
       reorthogonalize=True,

tensornetwork/matrixproductstates/dmrg.py

@@ -74,7 +74,7 @@ def __init__(self, mps: BaseMPS, mpo: BaseMPO, left_boundary: Tensor,
           'right_boundary.dtype = {} is different from BaseDMRG.dtype = {}'
           .format(self.right_envs[0].dtype, self.dtype))
 
-    def _single_site_matvec(L, mpotensor, R, mpstensor):
+    def _single_site_matvec(mpstensor, L, mpotensor, R):
       return ncon([L, mpstensor, mpotensor, R],
                   [[3, 1, -1], [1, 2, 4], [3, 5, -2, 2], [5, 4, -3]],
                   backend=self.backend.name)