How to use the netket.operator.Heisenberg function in NetKet

def _setup_model():
    g = nk.graph.Hypercube(8, 1)
    hi = nk.hilbert.Spin(g, 0.5)
    ham = nk.operator.Heisenberg(hi)

    ts = timestepper(hi.n_states, abs_tol=ATOL, rel_tol=RTOL)
    psi0 = np.random.rand(hi.n_states) + 1j * np.random.rand(hi.n_states)
    psi0 /= norm(psi0)

    return hi, ham, ts, psi0

def _setup():
    g = nk.graph.Hypercube(3, 2)
    hi = nk.hilbert.Spin(g, 0.5)

    ham = nk.operator.Heisenberg(hi)

    ma = nk.machine.RbmSpin(hi, alpha=2)
    ma.init_random_parameters()

    return hi, ham, ma

import netket as nk
import networkx as nx
import numpy as np

operators = {}

# Ising 1D
g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=0.5, graph=g)
operators["Ising 1D"] = nk.operator.Ising(h=1.321, hilbert=hi)

# Heisenberg 1D
g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=0.5, total_sz=0, graph=g)
operators["Heisenberg 1D"] = nk.operator.Heisenberg(hilbert=hi)

# Bose Hubbard
g = nk.graph.Hypercube(length=3, n_dim=2, pbc=True)
hi = nk.hilbert.Boson(n_max=3, n_bosons=6, graph=g)
operators["Bose Hubbard"] = nk.operator.BoseHubbard(U=4.0, hilbert=hi)

# Graph Hamiltonian
sigmax = [[0, 1], [1, 0]]
mszsz = [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
edges = [
    [0, 1],
    [1, 2],
    [2, 3],
    [3, 4],
    [4, 5],
    [5, 6],

# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import netket as nk

# 1D Lattice
g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
# with total Sz equal to 0
hi = nk.hilbert.Spin(s=0.5, graph=g, total_sz=0)

# Heisenberg hamiltonian
ha = nk.operator.Heisenberg(hilbert=hi)

# Symmetric RBM Spin Machine
ma = nk.machine.RbmSpinSymm(alpha=1, hilbert=hi)
ma.init_random_parameters(seed=1234, sigma=0.01)

# Metropolis Exchange Sampling
# Notice that this sampler exchanges two neighboring sites
# thus preservers the total magnetization
sa = nk.sampler.MetropolisExchange(machine=ma, graph=g)

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.05)

# Stochastic reconfiguration
gs = nk.variational.Vmc(
    hamiltonian=ha,

def run_netket(args):
    g = nk.graph.Hypercube(length=args.input_size, n_dim=1)
    hi = nk.hilbert.Spin(s=0.5, total_sz=0, graph=g)
    ha = nk.operator.Heisenberg(hilbert=hi)
    middle_layer = (nk.layer.ConvolutionalHypercube(length=args.input_size,
                                                    n_dim=1,
                                                    input_channels=args.width,
                                                    output_channels=args.width,
                                                    kernel_length=args.kernel_size),
                     nk.layer.Lncosh(input_size=args.width * args.input_size))
    middle_layers = middle_layer * (args.depth - 1)
    first_layer = (nk.layer.ConvolutionalHypercube(length=args.input_size,
                                                   n_dim=1,
                                                   input_channels=1,
                                                   output_channels=args.width,
                                                   kernel_length=args.kernel_size),
                 nk.layer.Lncosh(input_size=args.width * args.input_size),)
    ma = nk.machine.FFNN(hi, first_layer + middle_layers)
    ma.init_random_parameters(seed=1234, sigma=0.1)
    sa = nk.sampler.MetropolisHamiltonian(machine=ma, hamiltonian=ha)

# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import netket as nk

# 1D Lattice
g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
# with total Sz equal to 0
hi = nk.hilbert.Spin(s=0.5, graph=g, total_sz=0)

# Heisenberg hamiltonian
ha = nk.operator.Heisenberg(hilbert=hi)

# Symmetric RBM Spin Machine
ma = nk.machine.JastrowSymm(hilbert=hi)
ma.init_random_parameters(seed=1234, sigma=0.01)

# Metropolis Exchange Sampling
# Notice that this sampler exchanges two neighboring sites
# thus preservers the total magnetization
sa = nk.sampler.MetropolisExchange(machine=ma, graph=g)

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.05)

# Stochastic reconfiguration
gs = nk.variational.Vmc(
    hamiltonian=ha,

# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import netket as nk
from scipy.sparse.linalg import eigsh

# 1D Lattice
g = nk.graph.Hypercube(length=16, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
hi = nk.hilbert.Spin(s=0.5, graph=g)

# Heisenberg spin hamiltonian
ha = nk.operator.Heisenberg(hilbert=hi)

# Convert hamiltonian to a sparse matrix
# Here we further take only the real part since the Heisenberg Hamiltonian is real
sp_ha = ha.to_sparse().real

# Use scipy sparse diagonalization
vals, vecs = eigsh(sp_ha, k=2, which="SA")
print("eigenvalues with scipy sparse:", vals)

# Explicitely compute energy of ground state
# Doing full dot product
psi = vecs[:, 0]
print("\ng.s. energy:", psi @ sp_ha @ psi)

# Compute energy of first excited state
psi = vecs[:, 1]