Repository: netket/netket

Issue:
passing sparse matrices to LocalOperator
Hi,

so I would like to compute something like a structure factor for a large system. However I am limited by my Ram when defining all the matrices to pass to LocalOperator. Is it somehow possible to pass csr sparse matrices or something similiar? 

Thank you




Discussion:
Hi,

unfortunately this is not implemented yet,but it should be relatively
straightforward to modify LocalOperator in order to use local sparse
matrices. If there's a volunteer for that it would be great,otherwise I'll
look into this in September.

 Notice that the computational cost will stay proportional to the number of
nonzero elements in these natrices, so you might still be limited if your
local Hilbert space has a very large dimension and if you are considering
tensor products of several local operators

On Sat, Aug 1, 2020, 09:29 dambuck <notifications@github.com> wrote:

> Hi,
>
> so I would like to compute something like a structure factor for a large
> system. However I am limited by my Ram when defining all the matrices to
> pass to LocalOperator. Is it somehow possible to pass csr sparse matrices
> or something similiar?
>
> Thank you
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/netket/netket/issues/457>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AGWYRBBJQERE4OV6JSXRKFTR6QKFDANCNFSM4PRX3V6Q>
> .
>

Hi, thank you for the answer. 
Good to know
The problem I face here is, that I want to run netket through mpi, however, this also means that the operators I define are occupying N times the Ram, with N the number of processes.  This puts strong limitations on my Ram. Is there a way around this without loosing the parallel computation advantages? To somehow avoid saving N copies of the operators I need.

Thank you
Unfortunately this is not possible. You would need to completely rewrite the parallelization logic of netket using openmp instead of mpi... 
But again, I think that if LocalOperator matrices are too big, it means that most likely this is not the right object to work with to represent your hamiltonian. For example, we have specific cases covering bosonic operators with infinite local degrees of freedom or PauliStrings that cover cases that would take exp storage in the standard LocalOperator. 
> Hi, unfortunately this is not implemented yet,but it should be relatively straightforward to modify LocalOperator in order to use local sparse matrices. If there's a volunteer for that it would be great,otherwise I'll look into this in September. Notice that the computational cost will stay proportional to the number of nonzero elements in these natrices, so you might still be limited if your local Hilbert space has a very large dimension and if you are considering tensor products of several local operators
> […](#)
> On Sat, Aug 1, 2020, 09:29 dambuck ***@***.***> wrote: Hi, so I would like to compute something like a structure factor for a large system. However I am limited by my Ram when defining all the matrices to pass to LocalOperator. Is it somehow possible to pass csr sparse matrices or something similiar? Thank you — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#457>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AGWYRBBJQERE4OV6JSXRKFTR6QKFDANCNFSM4PRX3V6Q> .

This is not done yet, right @gcarleo ? I can do it
@gcarleo @PhilipVinc hi, i was wondering if this was taken care of. I will try to make very large 4 body boson operators and I'm afraid I will run on some memory issues / slow performance.
No.
There has literally been zero effort on this.

Part of the reason is that I plan to eventually rewrite most of the operator inner-working, especially of `LocalOperator`, to simplify the code and make it more jax-friendly.

`LocalOperator` has also grown to be a monstrosity that I'm not eager to touch, in fear of breaking something, as it's not tested extremely well.

You can try to do this, however. 
It might not be too complicated after all.
I think that replacing the calls to `np.array` or `np.asarray` used to convert the input matrix to `sp.sparse.coo_matrix` or `sp.sparse.csc_matrix` should work.
You should not touch the code dealing with constructing the inner representation, however.

Though, I fear it will be very slow.

Resolving patch:
diff --git a/netket/operator/_local_operator.py b/netket/operator/_local_operator.py
--- a/netket/operator/_local_operator.py
+++ b/netket/operator/_local_operator.py
@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import numbers
-from typing import Union, List, Optional
+from typing import Union, List, Optional, Any
 from netket.utils.types import DType, Array
 from textwrap import dedent
 
 import numpy as np
+import scipy.sparse as sp
 from numba import jit
 
 from netket.hilbert import AbstractHilbert, Fock
@@ -43,8 +45,10 @@ def _number_to_state(number, hilbert_size_per_site, local_states_per_site, out):
     return out
 
 
-def is_hermitian(a: np.ndarray, rtol=1e-05, atol=1e-08) -> bool:
-    return np.allclose(a, a.T.conj(), rtol=rtol, atol=atol)
+# TODO jit this. But scipy matrices are not jit-friendly so a solution must be found
+def is_hermitian(a: sp.csr_matrix, rtol=1e-05, atol=1e-08) -> bool:
+    c = np.abs(a - a.conj().T) - rtol * np.abs(a)
+    return c.max() <= atol
 
 
 def _dtype(obj: Union[numbers.Number, Array, "LocalOperator"]) -> DType:
@@ -247,13 +251,16 @@ def __init__(
             operators = [operators]
             acting_on = [acting_on]
 
-        operators = [np.asarray(operator) for operator in operators]
+        operators = [
+            operator if isinstance(operator, sp.csr_matrix) else sp.csr_matrix(operator)
+            for operator in operators
+        ]
 
         # If we asked for a specific dtype, enforce it.
         if dtype is None:
-            dtype = np.promote_types(operators[0].dtype, np.float32)
-            for op in operators[1:]:
-                np.promote_types(dtype, op.dtype)
+            dtype = functools.reduce(
+                lambda dt, op: np.promote_types(dt, op.dtype), operators, np.float32
+            )
 
         self._dtype = dtype
         self._init_zero()
@@ -271,9 +278,10 @@ def __init__(
     @property
     def operators(self) -> List[np.ndarray]:
         """List of the matrices of the operators encoded in this Local Operator.
-        Returns a copy.
+        Returns a copy of the operators.
         """
-        return self._operators_list()
+        operators = [op.copy().todense() for op in self._operators]
+        return operators
 
     @property
     def acting_on(self) -> List[List[int]]:
@@ -293,9 +301,14 @@ def size(self) -> int:
         return self._size
 
     @property
+    # A way to cache the property depending on modifications of self._operators is described here:
+    # https://stackoverflow.com/questions/48262273/python-bookkeeping-dependencies-in-cached-attributes-that-might-change
     def is_hermitian(self) -> bool:
         """Returns true if this operator is hermitian."""
-        return self._is_hermitian
+        # TODO: (VolodyaCO) I guess that if we have an operator with diagonal elements equal to 1j*C+Y, some complex constant, and
+        # self._constant=-1j*C, then the actual diagonal would be Y. How do we check hermiticity taking into account the diagonal
+        # elements as well as the self._constant? For the moment I just check hermiticity of the added constant, which must be real.
+        return np.all(self._is_hermitian_op) and np.isreal(self._constant)
 
     @property
     def mel_cutoff(self) -> float:
@@ -384,8 +397,9 @@ def __mul__(self, other):
         op._mels *= other
         op._constant *= other
 
-        for _op in op._operators:
+        for i, _op in enumerate(op._operators):
             _op *= other
+            op._is_hermitian_op[i] = is_hermitian(_op)
 
         op._nonzero_diagonal = has_nonzero_diagonal(op)
 
@@ -406,8 +420,9 @@ def __imul__(self, other):
         self._mels *= other
         self._constant *= other
 
-        for _op in self._operators:
+        for i, _op in enumerate(self._operators):
             _op *= other
+            self._is_hermitian_op[i] = is_hermitian(_op)
 
         self._nonzero_diagonal = has_nonzero_diagonal(self)
 
@@ -492,7 +507,9 @@ def _init_zero(self):
         self._local_states = np.zeros((0, 0, 0), dtype=np.float64)
 
         self._basis = np.zeros((0, 0), dtype=np.int64)
-        self._is_hermitian = True
+
+        # Array saving whether each operator is hermitian (plus one component to keep if self._constant is also hermitian).
+        self._is_hermitian_op = np.zeros((self._n_operators,), dtype=bool)
 
     def _acting_on_list(self):
         acting_on = []
@@ -501,17 +518,12 @@ def _acting_on_list(self):
 
         return acting_on
 
-    def _operators_list(self):
-        "A deep copy of the operators"
-        operators = [np.copy(op) for op in self._operators]
-        return operators
-
     def _add_operator(self, operator: Array, acting_on: List[int]):
         if not np.can_cast(operator, self.dtype, casting="same_kind"):
             raise ValueError(f"Cannot cast type {operator.dtype} to {self.dtype}")
 
         acting_on = np.asarray(acting_on, dtype=np.intp)
-        operator = np.asarray(operator, dtype=self.dtype)
+        operator = sp.csr_matrix(operator, dtype=self.dtype)
 
         if np.unique(acting_on).size != acting_on.size:
             raise ValueError("acting_on contains repeated entries.")
@@ -522,7 +534,7 @@ def _add_operator(self, operator: Array, acting_on: List[int]):
         if operator.ndim != 2:
             raise ValueError("The operator should be a matrix")
 
-        if np.all(np.abs(operator) < self.mel_cutoff):
+        if np.all(np.abs(operator.data) < self.mel_cutoff):
             return
 
         # re-sort the operator
@@ -540,15 +552,27 @@ def _add_operator(self, operator: Array, acting_on: List[int]):
         # If overlapping support, add the local operators themselves
         if support_i is not None:
             dim = min(operator.shape[0], self._operators[support_i].shape[0])
-            _opv = self._operators[support_i][:dim, :dim]
-            _opv += operator[:dim, :dim]
+            self._operators[support_i][:dim, :dim] += operator[:dim, :dim]
+            self._is_hermitian_op[support_i] = is_hermitian(self._operators[support_i])
 
             n_local_states_per_site = np.asarray(
                 [self.hilbert.size_at_index(i) for i in acting_on]
             )
-
+            assert (
+                self._operators[support_i].shape[0]
+                == self._operators[support_i].shape[1]
+            )
+            operator_row_index, operator_column_index = self._operators[
+                support_i
+            ].nonzero()
+            operator_data = np.asarray(
+                self._operators[support_i][operator_row_index, operator_column_index]
+            ).squeeze(axis=0)
             self._append_matrix(
-                self._operators[support_i],
+                operator_data,
+                operator_row_index,
+                operator_column_index,
+                self._operators[support_i].shape[0],
                 self._diag_mels[support_i],
                 self._mels[support_i],
                 self._x_prime[support_i],
@@ -558,17 +582,11 @@ def _add_operator(self, operator: Array, acting_on: List[int]):
                 self.mel_cutoff,
                 n_local_states_per_site,
             )
-
-            isherm = True
-            for op in self._operators:
-                isherm = isherm and is_hermitian(op)
-
-            self._is_hermitian = isherm
             self._nonzero_diagonal = has_nonzero_diagonal(self)
         else:
             self.__add_new_operator__(operator, acting_on)
 
-    def __add_new_operator__(self, operator: np.ndarray, acting_on: np.ndarray):
+    def __add_new_operator__(self, operator: sp.csr_matrix, acting_on: np.ndarray):
         # Else, we must add a completely new operator
         self._n_operators += 1
         self._operators.append(operator)
@@ -660,8 +678,21 @@ def __add_new_operator__(self, operator: np.ndarray, acting_on: np.ndarray):
         if acting_on.max() + 1 >= self._size:
             self._size = acting_on.max() + 1
 
+        self._is_hermitian_op = resize(
+            self._is_hermitian_op,
+            shape=(self.n_operators,),
+            init=is_hermitian(operator),
+        )
+
+        operator_row_index, operator_column_index = operator.nonzero()
+        operator_data = np.asarray(
+            operator[operator_row_index, operator_column_index]
+        ).squeeze(axis=0)
         self._append_matrix(
-            operator,
+            operator_data,
+            operator_row_index,
+            operator_column_index,
+            operator.shape[0],
             self._diag_mels[-1],
             self._mels[-1],
             self._x_prime[-1],
@@ -672,18 +703,15 @@ def __add_new_operator__(self, operator: np.ndarray, acting_on: np.ndarray):
             n_local_states_per_site,
         )
 
-        isherm = True
-        for op in self._operators:
-            isherm = isherm and is_hermitian(op)
-
-        self._is_hermitian = isherm
-
         self._nonzero_diagonal = has_nonzero_diagonal(self)
 
     @staticmethod
     @jit(nopython=True)
     def _append_matrix(
-        operator,
+        operator_data,
+        operator_row_index,
+        operator_column_index,
+        operator_size,
         diag_mels,
         mels,
         x_prime,
@@ -693,22 +721,41 @@ def _append_matrix(
         epsilon,
         hilb_size_per_site,
     ):
-        op_size = operator.shape[0]
-        assert op_size == operator.shape[1]
-        for i in range(op_size):
-            diag_mels[i] = operator[i, i]
-            n_conns[i] = 0
-            for j in range(op_size):
-                if i != j and np.abs(operator[i, j]) > epsilon:
-                    k_conn = n_conns[i]
-                    mels[i, k_conn] = operator[i, j]
+        n_conns[:operator_size] = 0
+        diag_mels[:operator_size] = 0
+        for element, ridx, cidx in zip(
+            operator_data, operator_row_index, operator_column_index
+        ):
+            if ridx == cidx:
+                diag_mels[ridx] = element
+            else:
+                if abs(element) > epsilon:
+                    k_conn = n_conns[ridx]
+                    mels[ridx, k_conn] = element
                     _number_to_state(
-                        j,
+                        cidx,
                         hilb_size_per_site,
                         local_states_per_site[:acting_size, :],
-                        x_prime[i, k_conn, :acting_size],
+                        x_prime[ridx, k_conn, :acting_size],
                     )
-                    n_conns[i] += 1
+                    n_conns[ridx] += 1
+        # CODE FOR COMPARISON:
+        # op_size = operator.shape[0]
+        # assert op_size == operator.shape[1]
+        # for i in range(op_size):
+        #     diag_mels[i] = operator[i, i]
+        #     n_conns[i] = 0
+        #     for j in range(op_size):
+        #         if i != j and np.abs(operator[i, j]) > epsilon:
+        #             k_conn = n_conns[i]
+        #             mels[i, k_conn] = operator[i, j]
+        #             _number_to_state(
+        #                 j,
+        #                 hilb_size_per_site,
+        #                 local_states_per_site[:acting_size, :],
+        #                 x_prime[i, k_conn, :acting_size],
+        #             )
+        #             n_conns[i] += 1
 
     def _multiply_operator(self, op, act):
         operators = []
@@ -722,11 +769,11 @@ def _multiply_operator(self, op, act):
 
             if act.size == act_i.size and np.array_equal(act, act_i):
                 # non-interesecting with same support
-                operators.append(np.copy(np.matmul(self._operators[i], op)))
+                operators.append(self._operators[i] @ op)
                 acting_on.append(act_i.tolist())
             elif inters.size == 0:
                 # disjoint supports
-                operators.append(np.copy(np.kron(self._operators[i], op)))
+                operators.append(sp.kron(self._operators[i], op, format="csr"))
                 acting_on.append(act_i.tolist() + act.tolist())
             else:
                 _act = list(act)
@@ -738,24 +785,28 @@ def _multiply_operator(self, op, act):
                 actmin = min(act)
                 for site in act_i:
                     if site not in act:
-                        I = np.eye(self.hilbert.shape[site], dtype=self.dtype)
+                        I = sp.eye(
+                            self.hilbert.shape[site], dtype=self.dtype, format="csr"
+                        )
                         if site < actmin:
                             _act = [site] + _act
-                            _op = np.kron(I, _op)
+                            _op = sp.kron(I, _op, format="csr")
                         else:  #  site > actmax
                             _act = _act + [site]
-                            _op = np.kron(_op, I)
+                            _op = sp.kron(_op, I, format="csr")
 
                 act_i_min = min(act_i)
                 for site in act:
                     if site not in act_i:
-                        I = np.eye(self.hilbert.shape[site], dtype=self.dtype)
+                        I = sp.eye(
+                            self.hilbert.shape[site], dtype=self.dtype, format="csr"
+                        )
                         if site < act_i_min:
                             _act_i = [site] + _act_i
-                            _op_i = np.kron(I, _op_i)
+                            _op_i = sp.kron(I, _op_i, format="csr")
                         else:  #  site > actmax
                             _act_i = _act_i + [site]
-                            _op_i = np.kron(_op_i, I)
+                            _op_i = sp.kron(_op_i, I, format="csr")
 
                 # reorder
                 _op, _act = _reorder_kronecker_product(self.hilbert, _op, _act)
@@ -763,14 +814,14 @@ def _multiply_operator(self, op, act):
 
                 if len(_act) == len(_act_i) and np.array_equal(_act, _act_i):
                     # non-interesecting with same support
-                    operators.append(np.matmul(_op_i, _op))
+                    operators.append(_op_i @ _op)
                     acting_on.append(_act_i)
                 else:
                     raise ValueError("Something failed")
 
         return operators, acting_on
 
-    def copy(self, *, dtype: Optional = None):
+    def copy(self, *, dtype: Optional[DType] = None):
         """Returns a copy of the operator, while optionally changing the dtype
         of the operator.
 
@@ -786,7 +837,7 @@ def copy(self, *, dtype: Optional = None):
 
         return LocalOperator(
             hilbert=self.hilbert,
-            operators=[np.copy(op) for op in self._operators],
+            operators=[op.copy() for op in self._operators],
             acting_on=self._acting_on_list(),
             constant=self._constant,
             dtype=dtype,
@@ -796,7 +847,7 @@ def transpose(self, *, concrete=False):
         r"""LocalOperator: Returns the tranpose of this operator."""
         if concrete:
 
-            new_ops = [np.copy(ops.transpose()) for ops in self._operators]
+            new_ops = [op.copy().transpose() for op in self._operators]
 
             return LocalOperator(
                 hilbert=self.hilbert,
@@ -809,7 +860,7 @@ def transpose(self, *, concrete=False):
 
     def conjugate(self, *, concrete=False):
         r"""LocalOperator: Returns the complex conjugate of this operator."""
-        new_ops = [np.copy(ops).conjugate() for ops in self._operators]
+        new_ops = [op.copy().conjugate() for op in self._operators]
 
         return LocalOperator(
             hilbert=self.hilbert,
@@ -823,9 +874,9 @@ def max_conn_size(self) -> int:
         """The maximum number of non zero ⟨x|O|x'⟩ for every x."""
         max_size = self.n_operators if self._nonzero_diagonal else 0
         for op in self._operators:
-            nnz_mat = np.abs(op) > self.mel_cutoff
-            nnz_mat[np.diag_indices(nnz_mat.shape[0])] = False
-            nnz_rows = np.sum(nnz_mat, axis=1)
+            nnz_rows = (abs(op) > self.mel_cutoff).getnnz(axis=1) - np.where(
+                np.abs(op.diagonal()) > 0, 1, 0
+            )
             max_size += np.max(nnz_rows)
 
         return max_size
diff --git a/netket/sampler/rules/custom_numpy.py b/netket/sampler/rules/custom_numpy.py
--- a/netket/sampler/rules/custom_numpy.py
+++ b/netket/sampler/rules/custom_numpy.py
@@ -129,10 +129,15 @@ def _choose_and_return(σp, x_prime, mels, sections, log_prob_corr, rnd_uniform)
     log_prob_corr.fill(0.0)
 
 
+def _sparse_allclose(a, b, rtol=1e-05, atol=1e-08):
+    c = np.abs(a - b) - rtol * np.abs(b)
+    return c.max() <= atol
+
+
 def _check_operators(operators):
     for op in operators:
         assert op.imag.max() < 1.0e-10
         assert op.min() >= 0
         assert np.allclose(op.sum(axis=0), 1.0)
         assert np.allclose(op.sum(axis=1), 1.0)
-        assert np.allclose(op, op.T)
+        assert _sparse_allclose(op, op.T)


Test patch:
diff --git a/test/operator/test_local_operator.py b/test/operator/test_local_operator.py
--- a/test/operator/test_local_operator.py
+++ b/test/operator/test_local_operator.py
@@ -318,7 +318,9 @@ def test_copy():
         assert op_copy is not op
         for o1, o2 in zip(op._operators, op_copy._operators):
             assert o1 is not o2
-            assert np.all(o1 == o2)
+            # Since operators might actually be sparse, accessing the dense data can be achieved by the .A attribute
+            # This also works for initially dense arrays.
+            assert np.all(o1.A == o2.A)
         same_matrices(op, op_copy)
 
 
@@ -328,6 +330,26 @@ def test_raises_unsorted_hilbert():
         nk.operator.LocalOperator(hi)
 
 
+def test_type_promotion():
+    hi = nk.hilbert.Qubit(1)
+    real_op = nk.operator.spin.sigmax(hi, 0, dtype=float)
+    complex_mat = nk.operator.spin.sigmay(hi, 0, dtype=complex).to_dense()
+    promoted_op = real_op + nk.operator.LocalOperator(hi, complex_mat, acting_on=[0])
+    assert promoted_op.dtype == np.complex128
+
+
+def test_is_hermitian():
+    for op in herm_operators.values():
+        assert op.is_hermitian == True
+
+    for (op, oph) in generic_operators.values():
+        assert op.is_hermitian == False
+        assert oph.is_hermitian == False
+
+    for op in herm_operators.values():
+        assert (1j * op).is_hermitian == False
+
+
 def test_qutip_conversion():
     # skip test if qutip not installed
     pytest.importorskip("qutip")