How to use NetKet - 10 common examples

def _test_stats_mean_std(hi, ham, ma, n_chains):
    sampler = nk.sampler.MetropolisLocal(ma, n_chains=n_chains)

    n_samples = 16000
    num_samples_per_chain = n_samples // n_chains

    # Discard a few samples
    sampler.generate_samples(1000)

    samples = sampler.generate_samples(num_samples_per_chain)
    assert samples.shape == (num_samples_per_chain, n_chains, hi.size)

    eloc = local_values(ham, ma, samples)
    assert eloc.shape == (num_samples_per_chain, n_chains)

    stats = statistics(eloc)

    # These tests only work for one MPI process

hilberts["Qubit"] = nk.hilbert.Qubit(graph=nk.graph.Hypercube(length=32, n_dim=1))

hilberts["Qubit n=1"] = nk.hilbert.Qubit()

# Custom Hilbert
hilberts["Custom Hilbert"] = nk.hilbert.CustomHilbert(
    local_states=[-1232, 132, 0], graph=nk.graph.Hypercube(length=34, n_dim=1)
)

hilberts["Custom Hilbert n"] = nk.hilbert.CustomHilbert(
    local_states=[-1232, 132, 0], n=2
)

# Heisenberg 1d
g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=0.5, total_sz=0.0, graph=g)
hilberts["Heisenberg 1d"] = hi

# Bose Hubbard
g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)
hi = nk.hilbert.Boson(n_max=4, n_bosons=20, graph=g)
hilberts["Bose Hubbard"] = hi

#
# Small hilbert space tests
#

# Spin 1/2
hilberts["Spin 1/2 Small"] = nk.hilbert.Spin(
    s=0.5, graph=nk.graph.Hypercube(length=10, n_dim=1)
)

hilberts["Spin 1/2"] = nk.hilbert.Spin(
    s=0.5, graph=nk.graph.Hypercube(length=20, n_dim=1)
)

hilberts["Spin 1/2 n"] = nk.hilbert.Spin(s=0.5, n=10)

# Spin 1/2 with total Sz
hilberts["Spin 1/2 with total Sz"] = nk.hilbert.Spin(
    s=0.5, total_sz=1.0, graph=nk.graph.Hypercube(length=20, n_dim=1)
)

# Spin 3
hilberts["Spin 3"] = nk.hilbert.Spin(s=3, graph=nk.graph.Hypercube(length=25, n_dim=1))

# Boson
hilberts["Boson"] = nk.hilbert.Boson(
    n_max=5, graph=nk.graph.Hypercube(length=21, n_dim=1)
)

# Boson with total number
hilberts["Bosons with total number"] = nk.hilbert.Boson(
    n_max=5, n_bosons=11, graph=nk.graph.Hypercube(length=21, n_dim=1)
)

# Qubit
hilberts["Qubit"] = nk.hilbert.Qubit(graph=nk.graph.Hypercube(length=32, n_dim=1))

hilberts["Qubit n=1"] = nk.hilbert.Qubit()

# Custom Hilbert
hilberts["Custom Hilbert"] = nk.hilbert.CustomHilbert(
    local_states=[-1232, 132, 0], graph=nk.graph.Hypercube(length=34, n_dim=1)

# Spin 1/2 with total Sz
hilberts["Spin 1/2 with total Sz"] = nk.hilbert.Spin(
    s=0.5, total_sz=1.0, graph=nk.graph.Hypercube(length=20, n_dim=1)
)

# Spin 3
hilberts["Spin 3"] = nk.hilbert.Spin(s=3, graph=nk.graph.Hypercube(length=25, n_dim=1))

# Boson
hilberts["Boson"] = nk.hilbert.Boson(
    n_max=5, graph=nk.graph.Hypercube(length=21, n_dim=1)
)

# Boson with total number
hilberts["Bosons with total number"] = nk.hilbert.Boson(
    n_max=5, n_bosons=11, graph=nk.graph.Hypercube(length=21, n_dim=1)
)

# Qubit
hilberts["Qubit"] = nk.hilbert.Qubit(graph=nk.graph.Hypercube(length=32, n_dim=1))

hilberts["Qubit n=1"] = nk.hilbert.Qubit()

# Custom Hilbert
hilberts["Custom Hilbert"] = nk.hilbert.CustomHilbert(
    local_states=[-1232, 132, 0], graph=nk.graph.Hypercube(length=34, n_dim=1)
)

hilberts["Custom Hilbert n"] = nk.hilbert.CustomHilbert(
    local_states=[-1232, 132, 0], n=2
)

import netket as nk
import networkx as nx
import igraph as ig
import math

nxg = nx.star_graph(10)
graphs = [
    nk.graph.Hypercube(length=10, n_dim=1, pbc=True),
    nk.graph.Hypercube(length=4, n_dim=2, pbc=True),
    nk.graph.Hypercube(length=5, n_dim=1, pbc=False),
    nk.graph.CustomGraph(nxg.edges()),
    nk.graph.Lattice(
        basis_vectors=[[1.0, 0.0], [1.0 / 2.0, math.sqrt(3) / 2.0]],
        extent=[10, 10],
        pbc=[0, 0],
        atoms_coord=[[0, 0]],
    ),
    nk.graph.Lattice(
        basis_vectors=[[1.5, math.sqrt(3) / 2.0], [0, math.sqrt(3)]],
        extent=[3, 5],
        atoms_coord=[[0, 0], [1, 0]],
    ),
    nk.graph.Lattice(
        basis_vectors=[[2.0, 0.0], [1.0, math.sqrt(3)]],
        extent=[4, 4],
        atoms_coord=[[0, 0], [1.0 / 2.0, math.sqrt(3) / 2.0], [1.0, 0.0]],

def _setup_vmc():
    g = nk.graph.Hypercube(length=8, n_dim=1)
    hi = nk.hilbert.Spin(s=0.5, graph=g)

    ma = nk.machine.RbmSpin(hilbert=hi, alpha=1)
    ma.init_random_parameters(seed=SEED, sigma=0.01)

    ha = nk.operator.Ising(hi, h=1.0)
    sa = nk.sampler.MetropolisLocal(machine=ma)
    sa.seed(SEED)
    op = nk.optimizer.Sgd(learning_rate=0.1)

    vmc = nk.variational.Vmc(
        hamiltonian=ha, sampler=sa, optimizer=op, n_samples=500, diag_shift=0.01
    )

    # Add custom observable
    X = [[0, 1], [1, 0]]

import netket as nk
import networkx as nx
import numpy as np
import pytest
from pytest import approx

samplers = {}

# TESTS FOR SPIN HILBERT
# Constructing a 1d lattice
g = nk.graph.Hypercube(length=6, n_dim=1)

# Hilbert space of spins from given graph
hi = nk.hilbert.Spin(s=0.5, graph=g)
ma = nk.machine.RbmSpin(hilbert=hi, alpha=1)
ma.init_random_parameters(seed=1234, sigma=0.2)

sa = nk.sampler.MetropolisLocal(machine=ma)
samplers["MetropolisLocal RbmSpin"] = sa

sa = nk.sampler.MetropolisLocalPt(machine=ma, n_replicas=4)
samplers["MetropolisLocalPt RbmSpin"] = sa

ha = nk.operator.Ising(hilbert=hi, h=1.0)
sa = nk.sampler.MetropolisHamiltonian(machine=ma, hamiltonian=ha)
samplers["MetropolisHamiltonian RbmSpin"] = sa

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

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

def test_no_segfault():
    g = nk.graph.Hypercube(8, 1)
    hi = nk.hilbert.Spin(g, 0.5)

    lo = nk.operator.LocalOperator(hi, [[1, 0], [0, 1]], [0])
    lo = lo.transpose()

    hi = None

    lo = lo * lo

    assert True