We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Bose-Hubbard model fails in extended Hilbert space for newer versions
Hello, 
we are working on an extended Hilbert space in order to calculate the ground state energy of the Bose-Hubbard model. We added spin 1/2 sites between every boson site in the one-dimensional chain, but we left the Hamiltonian unchanged. We expected that the ground state energy should be the same in this extended Hilbert space, but in the newer versions of Netket (post 3.3.3), this energy differs from the actual solution. Even more, every time we run the exact diagonalization of the Hamiltonian matrix (either with the full_ed or lanczos_ed method), its eigenvalues change; we also diagonalized the same Hamiltonian matrix with the linalg package of numpy and the result also differs every time we do it. This doesn't happen in the older versions; the problem begins in the 3.3.4 version of Netket.

We claim that the problem resides in the way that the newer versions of Netket interpret the Hamiltonian matrices with mixed Hilbert spaces. When we do the same Bose-Hubbard model in a Fock-Boson space, the calculation gives the correct value of the ground state energy.

Here is the code that we run in both versions of Netket, and a video of the discrepancy. 

Notebook ran in Nektet 3.3.3: https://www.youtube.com/watch?v=ENhRJfYg7dg
Notebook ran in Netket 3.4.1: https://www.youtube.com/watch?v=Q3XfWrnR7LU

You can download the Jupyter Script here:
https://www.dropbox.com/s/dza4kbyem2ycg6v/BoseHubbardNetket.ipynb?dl=0

```
# Extended Hilbert space
n_max = 3
L = 3 # Number of spin sites
hil_gen = nk.hilbert.Fock(n_max = n_max)*nk.hilbert.Spin(1/2)
hi_ext = hil_gen
for i in range(L-1):
    hi_ext = hi_ext * hil_gen 
hi_ext = hi_ext * nk.hilbert.Fock(n_max = n_max)
# Boson-Fock Hilbert space
N = 4 # Number of bosons in the chain
# Chain graph
g = nk.graph.Chain(length=N, pbc=False)
hi_fock = nk.hilbert.Fock(n_max=n_max, N=N)

# Bose-Hubbard Hamiltonian for extended Hilbert space
J = 0.5
U = 0.5
h = create(hi_ext,0)*destroy(hi_ext,2) + create(hi_ext,2)*destroy(hi_ext,4) + create(hi_ext,4)*destroy(hi_ext,6)
h_hc = h.H.collect()
h_u = (number(hi_ext,0)*(number(hi_ext,0)-1) + number(hi_ext,2)*(number(hi_ext,2)-1) + number(hi_ext,4)*(number(hi_ext,4)-1) + 
      number(hi_ext,6)*(number(hi_ext,6)-1))
H_bh = -J*(h + h_hc) + (U/2)*h_u

# Bose-Hubbard Hamiltonian for the Fock-Bose Hilbert space
H_bh2 = nk.operator.BoseHubbard(hilbert=hi_fock,U=0.5,J=0.5,graph=g)

# lanczos computation of the eigenvalues of the Hamiltonian with the extended Hilbert space
E_bh, ket_bh = nk.exact.lanczos_ed(H_bh, compute_eigenvectors=True)
print("Exact ground state energy = {0:.3f}".format(E_bh[0]))

# full computation of the eigenvalues of the Hamiltonian with the extended Hilbert space
E_bhfull = nk.exact.full_ed(H_bh)
print("Exact ground state energy = {0:.3f}".format(E_bhfull[0]))

# numpy.linalg computation of the eigenvalues of the Hamiltonian with the extended Hilbert space
M_bh = H_bh.to_dense()
E_bhla = min(np.linalg.eig(M_bh)[0])
print("Ground state energy:", E_bhla.real)

# scipy sparse computation of the eigenvalues of the Hamiltonian with the extended Hilbert space
H_bhsparse = H_bh.to_sparse()
eig_vals, eig_vecs = eigsh(H_bhsparse, k=1, which="SA")
print("eigenvalues with scipy sparse:", eig_vals)

# lanczos computation of the eigenvalues of the Hamiltonian with the Fock Hilbert space
E_bh2, ket_bh2 = nk.exact.lanczos_ed(H_bh2, compute_eigenvectors=True)
print("Exact ground state energy = {0:.3f}".format(E_bh2[0]))

# full computation of the eigenvalues of the Hamiltonian with the Fock Hilbert space
E_bhfull2= nk.exact.full_ed(H_bh2)
print("Exact ground state energy = {0:.3f}".format(E_bhfull2[0]))

# numpy.linalg computation of the eigenvalues of the Hamiltonian with the Fock Hilbert space
M_bh2 = H_bh2.to_dense()
E_bhla2 = min(np.linalg.eig(M_bh2)[0])
print("Ground state energy:", E_bhla2.real)

# scipy sparse computation of the eigenvalues of the Hamiltonian with the extended Hilbert space
H_bhsparse2 = H_bh2.to_sparse()
eig_vals2, eig_vecs2 = eigsh(H_bhsparse2, k=1, which="SA")
print("eigenvalues with scipy sparse:", eig_vals2)
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `netket/operator/_local_operator_compile_helpers.py`
Content:
```
1 # Copyright 2022 The NetKet Authors - All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #    http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """
16 This file contains functions generating the numba-packed representation of local
17 operators.
18 """
19 
20 import numpy as np
21 import numba
22 
23 from netket.hilbert import AbstractHilbert
24 from netket.utils.types import DType
25 
26 
27 def pack_internals(
28     hilbert: AbstractHilbert,
29     operators_dict: dict,
30     constant,
31     dtype: DType,
32     mel_cutoff: float,
33 ):
34     """
35     Take the internal lazy representation of a local operator and returns the arrays
36     needed for the numba implementation.
37 
38     This takes as input a dictionary with Tuples as keys, the `acting_on` and matrices as values.
39     The keys represent the sites upon which the matrix acts.
40     It is assumed that the integer in the tuples are sorted.
41 
42     Returns a dictionary with all the data fields
43     """
44     op_acting_on = list(operators_dict.keys())
45     operators = list(operators_dict.values())
46     n_operators = len(operators_dict)
47 
48     """Analyze the operator strings and precompute arrays for get_conn inference"""
49     acting_size = np.array([len(aon) for aon in op_acting_on], dtype=np.intp)
50     max_acting_on_sz = np.max(acting_size)
51     max_local_hilbert_size = max(
52         [max(map(hilbert.size_at_index, aon)) for aon in op_acting_on]
53     )
54     max_op_size = max(map(lambda x: x.shape[0], operators))
55 
56     acting_on = np.full((n_operators, max_acting_on_sz), -1, dtype=np.intp)
57     for (i, aon) in enumerate(op_acting_on):
58         acting_on[i][: len(aon)] = aon
59 
60     local_states = np.full(
61         (n_operators, max_acting_on_sz, max_local_hilbert_size), np.nan
62     )
63     basis = np.full((n_operators, max_acting_on_sz), 1e10, dtype=np.int64)
64 
65     diag_mels = np.full((n_operators, max_op_size), np.nan, dtype=dtype)
66     mels = np.full(
67         (n_operators, max_op_size, max_op_size - 1),
68         np.nan,
69         dtype=dtype,
70     )
71     x_prime = np.full(
72         (n_operators, max_op_size, max_op_size - 1, max_acting_on_sz),
73         -1,
74         dtype=np.float64,
75     )
76     n_conns = np.full((n_operators, max_op_size), -1, dtype=np.intp)
77 
78     for (i, (aon, op)) in enumerate(operators_dict.items()):
79         aon_size = len(aon)
80         n_local_states_per_site = np.asarray([hilbert.size_at_index(i) for i in aon])
81 
82         ## add an operator to local_states
83         for (j, site) in enumerate(aon):
84             local_states[i, j, : hilbert.shape[site]] = np.asarray(
85                 hilbert.states_at_index(site)
86             )
87 
88         ba = 1
89         for s in range(aon_size):
90             basis[i, s] = ba
91             ba *= hilbert.shape[aon_size - s - 1]
92 
93         # eventually could support sparse matrices
94         # if isinstance(op, sparse.spmatrix):
95         #    op = op.todense()
96 
97         _append_matrix(
98             op,
99             diag_mels[i],
100             mels[i],
101             x_prime[i],
102             n_conns[i],
103             aon_size,
104             local_states[i],
105             mel_cutoff,
106             n_local_states_per_site,
107         )
108 
109     nonzero_diagonal = (
110         np.any(np.abs(diag_mels) >= mel_cutoff) or np.abs(constant) >= mel_cutoff
111     )
112 
113     max_conn_size = 1 if nonzero_diagonal else 0
114     for op in operators:
115         nnz_mat = np.abs(op) > mel_cutoff
116         nnz_mat[np.diag_indices(nnz_mat.shape[0])] = False
117         nnz_rows = np.sum(nnz_mat, axis=1)
118         max_conn_size += np.max(nnz_rows)
119 
120     return {
121         "acting_on": acting_on,
122         "acting_size": acting_size,
123         "diag_mels": diag_mels,
124         "mels": mels,
125         "x_prime": x_prime,
126         "n_conns": n_conns,
127         "local_states": local_states,
128         "basis": basis,
129         "nonzero_diagonal": nonzero_diagonal,
130         "max_conn_size": max_conn_size,
131     }
132 
133 
134 @numba.jit(nopython=True)
135 def _append_matrix(
136     operator,
137     diag_mels,
138     mels,
139     x_prime,
140     n_conns,
141     acting_size,
142     local_states_per_site,
143     epsilon,
144     hilb_size_per_site,
145 ):
146     op_size = operator.shape[0]
147     assert op_size == operator.shape[1]
148     for i in range(op_size):
149         diag_mels[i] = operator[i, i]
150         n_conns[i] = 0
151         for j in range(op_size):
152             if i != j and np.abs(operator[i, j]) > epsilon:
153                 k_conn = n_conns[i]
154                 mels[i, k_conn] = operator[i, j]
155                 _number_to_state(
156                     j,
157                     hilb_size_per_site,
158                     local_states_per_site[:acting_size, :],
159                     x_prime[i, k_conn, :acting_size],
160                 )
161                 n_conns[i] += 1
162 
163 
164 @numba.jit(nopython=True)
165 def _number_to_state(number, hilbert_size_per_site, local_states_per_site, out):
166 
167     out[:] = local_states_per_site[:, 0]
168     size = out.shape[0]
169 
170     ip = number
171     k = size - 1
172     while ip > 0:
173         local_size = hilbert_size_per_site[k]
174         out[k] = local_states_per_site[k, ip % local_size]
175         ip = ip // local_size
176         k -= 1
177 
178     return out
179
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/netket/operator/_local_operator_compile_helpers.py b/netket/operator/_local_operator_compile_helpers.py
--- a/netket/operator/_local_operator_compile_helpers.py
+++ b/netket/operator/_local_operator_compile_helpers.py
@@ -88,7 +88,7 @@
         ba = 1
         for s in range(aon_size):
             basis[i, s] = ba
-            ba *= hilbert.shape[aon_size - s - 1]
+            ba *= hilbert.shape[aon[aon_size - s - 1]]
 
         # eventually could support sparse matrices
         # if isinstance(op, sparse.spmatrix):

{"golden_diff": "diff --git a/netket/operator/_local_operator_compile_helpers.py b/netket/operator/_local_operator_compile_helpers.py\n--- a/netket/operator/_local_operator_compile_helpers.py\n+++ b/netket/operator/_local_operator_compile_helpers.py\n@@ -88,7 +88,7 @@\n         ba = 1\n         for s in range(aon_size):\n             basis[i, s] = ba\n-            ba *= hilbert.shape[aon_size - s - 1]\n+            ba *= hilbert.shape[aon[aon_size - s - 1]]\n \n         # eventually could support sparse matrices\n         # if isinstance(op, sparse.spmatrix):\n", "issue": "Bose-Hubbard model fails in extended Hilbert space for newer versions\nHello, \r\nwe are working on an extended Hilbert space in order to calculate the ground state energy of the Bose-Hubbard model. We added spin 1/2 sites between every boson site in the one-dimensional chain, but we left the Hamiltonian unchanged. We expected that the ground state energy should be the same in this extended Hilbert space, but in the newer versions of Netket (post 3.3.3), this energy differs from the actual solution. Even more, every time we run the exact diagonalization of the Hamiltonian matrix (either with the full_ed or lanczos_ed method), its eigenvalues change; we also diagonalized the same Hamiltonian matrix with the linalg package of numpy and the result also differs every time we do it. This doesn't happen in the older versions; the problem begins in the 3.3.4 version of Netket.\r\n\r\nWe claim that the problem resides in the way that the newer versions of Netket interpret the Hamiltonian matrices with mixed Hilbert spaces. When we do the same Bose-Hubbard model in a Fock-Boson space, the calculation gives the correct value of the ground state energy.\r\n\r\nHere is the code that we run in both versions of Netket, and a video of the discrepancy. \r\n\r\nNotebook ran in Nektet 3.3.3: https://www.youtube.com/watch?v=ENhRJfYg7dg\r\nNotebook ran in Netket 3.4.1: https://www.youtube.com/watch?v=Q3XfWrnR7LU\r\n\r\nYou can download the Jupyter Script here:\r\nhttps://www.dropbox.com/s/dza4kbyem2ycg6v/BoseHubbardNetket.ipynb?dl=0\r\n\r\n```\r\n# Extended Hilbert space\r\nn_max = 3\r\nL = 3 # Number of spin sites\r\nhil_gen = nk.hilbert.Fock(n_max = n_max)*nk.hilbert.Spin(1/2)\r\nhi_ext = hil_gen\r\nfor i in range(L-1):\r\n    hi_ext = hi_ext * hil_gen \r\nhi_ext = hi_ext * nk.hilbert.Fock(n_max = n_max)\r\n# Boson-Fock Hilbert space\r\nN = 4 # Number of bosons in the chain\r\n# Chain graph\r\ng = nk.graph.Chain(length=N, pbc=False)\r\nhi_fock = nk.hilbert.Fock(n_max=n_max, N=N)\r\n\r\n# Bose-Hubbard Hamiltonian for extended Hilbert space\r\nJ = 0.5\r\nU = 0.5\r\nh = create(hi_ext,0)*destroy(hi_ext,2) + create(hi_ext,2)*destroy(hi_ext,4) + create(hi_ext,4)*destroy(hi_ext,6)\r\nh_hc = h.H.collect()\r\nh_u = (number(hi_ext,0)*(number(hi_ext,0)-1) + number(hi_ext,2)*(number(hi_ext,2)-1) + number(hi_ext,4)*(number(hi_ext,4)-1) + \r\n      number(hi_ext,6)*(number(hi_ext,6)-1))\r\nH_bh = -J*(h + h_hc) + (U/2)*h_u\r\n\r\n# Bose-Hubbard Hamiltonian for the Fock-Bose Hilbert space\r\nH_bh2 = nk.operator.BoseHubbard(hilbert=hi_fock,U=0.5,J=0.5,graph=g)\r\n\r\n# lanczos computation of the eigenvalues of the Hamiltonian with the extended Hilbert space\r\nE_bh, ket_bh = nk.exact.lanczos_ed(H_bh, compute_eigenvectors=True)\r\nprint(\"Exact ground state energy = {0:.3f}\".format(E_bh[0]))\r\n\r\n# full computation of the eigenvalues of the Hamiltonian with the extended Hilbert space\r\nE_bhfull = nk.exact.full_ed(H_bh)\r\nprint(\"Exact ground state energy = {0:.3f}\".format(E_bhfull[0]))\r\n\r\n# numpy.linalg computation of the eigenvalues of the Hamiltonian with the extended Hilbert space\r\nM_bh = H_bh.to_dense()\r\nE_bhla = min(np.linalg.eig(M_bh)[0])\r\nprint(\"Ground state energy:\", E_bhla.real)\r\n\r\n# scipy sparse computation of the eigenvalues of the Hamiltonian with the extended Hilbert space\r\nH_bhsparse = H_bh.to_sparse()\r\neig_vals, eig_vecs = eigsh(H_bhsparse, k=1, which=\"SA\")\r\nprint(\"eigenvalues with scipy sparse:\", eig_vals)\r\n\r\n# lanczos computation of the eigenvalues of the Hamiltonian with the Fock Hilbert space\r\nE_bh2, ket_bh2 = nk.exact.lanczos_ed(H_bh2, compute_eigenvectors=True)\r\nprint(\"Exact ground state energy = {0:.3f}\".format(E_bh2[0]))\r\n\r\n# full computation of the eigenvalues of the Hamiltonian with the Fock Hilbert space\r\nE_bhfull2= nk.exact.full_ed(H_bh2)\r\nprint(\"Exact ground state energy = {0:.3f}\".format(E_bhfull2[0]))\r\n\r\n# numpy.linalg computation of the eigenvalues of the Hamiltonian with the Fock Hilbert space\r\nM_bh2 = H_bh2.to_dense()\r\nE_bhla2 = min(np.linalg.eig(M_bh2)[0])\r\nprint(\"Ground state energy:\", E_bhla2.real)\r\n\r\n# scipy sparse computation of the eigenvalues of the Hamiltonian with the extended Hilbert space\r\nH_bhsparse2 = H_bh2.to_sparse()\r\neig_vals2, eig_vecs2 = eigsh(H_bhsparse2, k=1, which=\"SA\")\r\nprint(\"eigenvalues with scipy sparse:\", eig_vals2)\r\n```\n", "before_files": [{"content": "# Copyright 2022 The NetKet Authors - All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis file contains functions generating the numba-packed representation of local\noperators.\n\"\"\"\n\nimport numpy as np\nimport numba\n\nfrom netket.hilbert import AbstractHilbert\nfrom netket.utils.types import DType\n\n\ndef pack_internals(\n    hilbert: AbstractHilbert,\n    operators_dict: dict,\n    constant,\n    dtype: DType,\n    mel_cutoff: float,\n):\n    \"\"\"\n    Take the internal lazy representation of a local operator and returns the arrays\n    needed for the numba implementation.\n\n    This takes as input a dictionary with Tuples as keys, the `acting_on` and matrices as values.\n    The keys represent the sites upon which the matrix acts.\n    It is assumed that the integer in the tuples are sorted.\n\n    Returns a dictionary with all the data fields\n    \"\"\"\n    op_acting_on = list(operators_dict.keys())\n    operators = list(operators_dict.values())\n    n_operators = len(operators_dict)\n\n    \"\"\"Analyze the operator strings and precompute arrays for get_conn inference\"\"\"\n    acting_size = np.array([len(aon) for aon in op_acting_on], dtype=np.intp)\n    max_acting_on_sz = np.max(acting_size)\n    max_local_hilbert_size = max(\n        [max(map(hilbert.size_at_index, aon)) for aon in op_acting_on]\n    )\n    max_op_size = max(map(lambda x: x.shape[0], operators))\n\n    acting_on = np.full((n_operators, max_acting_on_sz), -1, dtype=np.intp)\n    for (i, aon) in enumerate(op_acting_on):\n        acting_on[i][: len(aon)] = aon\n\n    local_states = np.full(\n        (n_operators, max_acting_on_sz, max_local_hilbert_size), np.nan\n    )\n    basis = np.full((n_operators, max_acting_on_sz), 1e10, dtype=np.int64)\n\n    diag_mels = np.full((n_operators, max_op_size), np.nan, dtype=dtype)\n    mels = np.full(\n        (n_operators, max_op_size, max_op_size - 1),\n        np.nan,\n        dtype=dtype,\n    )\n    x_prime = np.full(\n        (n_operators, max_op_size, max_op_size - 1, max_acting_on_sz),\n        -1,\n        dtype=np.float64,\n    )\n    n_conns = np.full((n_operators, max_op_size), -1, dtype=np.intp)\n\n    for (i, (aon, op)) in enumerate(operators_dict.items()):\n        aon_size = len(aon)\n        n_local_states_per_site = np.asarray([hilbert.size_at_index(i) for i in aon])\n\n        ## add an operator to local_states\n        for (j, site) in enumerate(aon):\n            local_states[i, j, : hilbert.shape[site]] = np.asarray(\n                hilbert.states_at_index(site)\n            )\n\n        ba = 1\n        for s in range(aon_size):\n            basis[i, s] = ba\n            ba *= hilbert.shape[aon_size - s - 1]\n\n        # eventually could support sparse matrices\n        # if isinstance(op, sparse.spmatrix):\n        #    op = op.todense()\n\n        _append_matrix(\n            op,\n            diag_mels[i],\n            mels[i],\n            x_prime[i],\n            n_conns[i],\n            aon_size,\n            local_states[i],\n            mel_cutoff,\n            n_local_states_per_site,\n        )\n\n    nonzero_diagonal = (\n        np.any(np.abs(diag_mels) >= mel_cutoff) or np.abs(constant) >= mel_cutoff\n    )\n\n    max_conn_size = 1 if nonzero_diagonal else 0\n    for op in operators:\n        nnz_mat = np.abs(op) > mel_cutoff\n        nnz_mat[np.diag_indices(nnz_mat.shape[0])] = False\n        nnz_rows = np.sum(nnz_mat, axis=1)\n        max_conn_size += np.max(nnz_rows)\n\n    return {\n        \"acting_on\": acting_on,\n        \"acting_size\": acting_size,\n        \"diag_mels\": diag_mels,\n        \"mels\": mels,\n        \"x_prime\": x_prime,\n        \"n_conns\": n_conns,\n        \"local_states\": local_states,\n        \"basis\": basis,\n        \"nonzero_diagonal\": nonzero_diagonal,\n        \"max_conn_size\": max_conn_size,\n    }\n\n\n@numba.jit(nopython=True)\ndef _append_matrix(\n    operator,\n    diag_mels,\n    mels,\n    x_prime,\n    n_conns,\n    acting_size,\n    local_states_per_site,\n    epsilon,\n    hilb_size_per_site,\n):\n    op_size = operator.shape[0]\n    assert op_size == operator.shape[1]\n    for i in range(op_size):\n        diag_mels[i] = operator[i, i]\n        n_conns[i] = 0\n        for j in range(op_size):\n            if i != j and np.abs(operator[i, j]) > epsilon:\n                k_conn = n_conns[i]\n                mels[i, k_conn] = operator[i, j]\n                _number_to_state(\n                    j,\n                    hilb_size_per_site,\n                    local_states_per_site[:acting_size, :],\n                    x_prime[i, k_conn, :acting_size],\n                )\n                n_conns[i] += 1\n\n\n@numba.jit(nopython=True)\ndef _number_to_state(number, hilbert_size_per_site, local_states_per_site, out):\n\n    out[:] = local_states_per_site[:, 0]\n    size = out.shape[0]\n\n    ip = number\n    k = size - 1\n    while ip > 0:\n        local_size = hilbert_size_per_site[k]\n        out[k] = local_states_per_site[k, ip % local_size]\n        ip = ip // local_size\n        k -= 1\n\n    return out\n", "path": "netket/operator/_local_operator_compile_helpers.py"}], "after_files": [{"content": "# Copyright 2022 The NetKet Authors - All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis file contains functions generating the numba-packed representation of local\noperators.\n\"\"\"\n\nimport numpy as np\nimport numba\n\nfrom netket.hilbert import AbstractHilbert\nfrom netket.utils.types import DType\n\n\ndef pack_internals(\n    hilbert: AbstractHilbert,\n    operators_dict: dict,\n    constant,\n    dtype: DType,\n    mel_cutoff: float,\n):\n    \"\"\"\n    Take the internal lazy representation of a local operator and returns the arrays\n    needed for the numba implementation.\n\n    This takes as input a dictionary with Tuples as keys, the `acting_on` and matrices as values.\n    The keys represent the sites upon which the matrix acts.\n    It is assumed that the integer in the tuples are sorted.\n\n    Returns a dictionary with all the data fields\n    \"\"\"\n    op_acting_on = list(operators_dict.keys())\n    operators = list(operators_dict.values())\n    n_operators = len(operators_dict)\n\n    \"\"\"Analyze the operator strings and precompute arrays for get_conn inference\"\"\"\n    acting_size = np.array([len(aon) for aon in op_acting_on], dtype=np.intp)\n    max_acting_on_sz = np.max(acting_size)\n    max_local_hilbert_size = max(\n        [max(map(hilbert.size_at_index, aon)) for aon in op_acting_on]\n    )\n    max_op_size = max(map(lambda x: x.shape[0], operators))\n\n    acting_on = np.full((n_operators, max_acting_on_sz), -1, dtype=np.intp)\n    for (i, aon) in enumerate(op_acting_on):\n        acting_on[i][: len(aon)] = aon\n\n    local_states = np.full(\n        (n_operators, max_acting_on_sz, max_local_hilbert_size), np.nan\n    )\n    basis = np.full((n_operators, max_acting_on_sz), 1e10, dtype=np.int64)\n\n    diag_mels = np.full((n_operators, max_op_size), np.nan, dtype=dtype)\n    mels = np.full(\n        (n_operators, max_op_size, max_op_size - 1),\n        np.nan,\n        dtype=dtype,\n    )\n    x_prime = np.full(\n        (n_operators, max_op_size, max_op_size - 1, max_acting_on_sz),\n        -1,\n        dtype=np.float64,\n    )\n    n_conns = np.full((n_operators, max_op_size), -1, dtype=np.intp)\n\n    for (i, (aon, op)) in enumerate(operators_dict.items()):\n        aon_size = len(aon)\n        n_local_states_per_site = np.asarray([hilbert.size_at_index(i) for i in aon])\n\n        ## add an operator to local_states\n        for (j, site) in enumerate(aon):\n            local_states[i, j, : hilbert.shape[site]] = np.asarray(\n                hilbert.states_at_index(site)\n            )\n\n        ba = 1\n        for s in range(aon_size):\n            basis[i, s] = ba\n            ba *= hilbert.shape[aon[aon_size - s - 1]]\n\n        # eventually could support sparse matrices\n        # if isinstance(op, sparse.spmatrix):\n        #    op = op.todense()\n\n        _append_matrix(\n            op,\n            diag_mels[i],\n            mels[i],\n            x_prime[i],\n            n_conns[i],\n            aon_size,\n            local_states[i],\n            mel_cutoff,\n            n_local_states_per_site,\n        )\n\n    nonzero_diagonal = (\n        np.any(np.abs(diag_mels) >= mel_cutoff) or np.abs(constant) >= mel_cutoff\n    )\n\n    max_conn_size = 1 if nonzero_diagonal else 0\n    for op in operators:\n        nnz_mat = np.abs(op) > mel_cutoff\n        nnz_mat[np.diag_indices(nnz_mat.shape[0])] = False\n        nnz_rows = np.sum(nnz_mat, axis=1)\n        max_conn_size += np.max(nnz_rows)\n\n    return {\n        \"acting_on\": acting_on,\n        \"acting_size\": acting_size,\n        \"diag_mels\": diag_mels,\n        \"mels\": mels,\n        \"x_prime\": x_prime,\n        \"n_conns\": n_conns,\n        \"local_states\": local_states,\n        \"basis\": basis,\n        \"nonzero_diagonal\": nonzero_diagonal,\n        \"max_conn_size\": max_conn_size,\n    }\n\n\n@numba.jit(nopython=True)\ndef _append_matrix(\n    operator,\n    diag_mels,\n    mels,\n    x_prime,\n    n_conns,\n    acting_size,\n    local_states_per_site,\n    epsilon,\n    hilb_size_per_site,\n):\n    op_size = operator.shape[0]\n    assert op_size == operator.shape[1]\n    for i in range(op_size):\n        diag_mels[i] = operator[i, i]\n        n_conns[i] = 0\n        for j in range(op_size):\n            if i != j and np.abs(operator[i, j]) > epsilon:\n                k_conn = n_conns[i]\n                mels[i, k_conn] = operator[i, j]\n                _number_to_state(\n                    j,\n                    hilb_size_per_site,\n                    local_states_per_site[:acting_size, :],\n                    x_prime[i, k_conn, :acting_size],\n                )\n                n_conns[i] += 1\n\n\n@numba.jit(nopython=True)\ndef _number_to_state(number, hilbert_size_per_site, local_states_per_site, out):\n\n    out[:] = local_states_per_site[:, 0]\n    size = out.shape[0]\n\n    ip = number\n    k = size - 1\n    while ip > 0:\n        local_size = hilbert_size_per_site[k]\n        out[k] = local_states_per_site[k, ip % local_size]\n        ip = ip // local_size\n        k -= 1\n\n    return out\n", "path": "netket/operator/_local_operator_compile_helpers.py"}]}