How to use the netket.sampler.MetropolisLocal function in NetKet

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

# Constructing a 1d lattice
g = nk.graph.Hypercube(length=20, n_dim=1)

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

# Hamiltonian
ha = nk.operator.Ising(h=1.0, hilbert=hi)

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

# Sampler
sa = nk.sampler.MetropolisLocal(machine=ma)
sa.seed(SEED)

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

# Variational Monte Carlo
vmc = nk.variational.Vmc(
    hamiltonian=ha,
    sampler=sa,
    optimizer=op,
    n_samples=1000,
    diag_shift=0.0,
    method="Sr",
)

mpi_rank = nk.MPI.rank()

# Hypercube
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

# Custom Hilbert Space
hi = nk.hilbert.Spin(s=0.5, graph=g)

# Graph Operator
op = nk.operator.GraphOperator(hi, siteops=site_operator, bondops=bond_operator)

# Restricted Boltzmann Machine
ma = nk.machine.RbmSpin(hilbert=hi, alpha=1)
ma.init_random_parameters(seed=1234, sigma=0.01)

# Local Metropolis Sampling
sa = nk.sampler.MetropolisLocal(machine=ma)

# Optimizer
opt = nk.optimizer.AdaMax()

# Stochastic reconfiguration
gs = nk.variational.Vmc(
    hamiltonian=op,
    sampler=sa,
    optimizer=opt,
    n_samples=1000,
    diag_shift=0.1,
    method="Gd",
)

gs.run(output_prefix="test", n_iter=30000)

ma = nk.machine.RbmMultiVal(hilbert=hi, alpha=10)

## Layers
#L = 28*28
#layers = (
#    nk.layer.FullyConnected(input_size=L, output_size=100),
#    nk.layer.Relu(input_size=100),
#    nk.layer.FullyConnected(input_size=100, output_size=10),
#)
#
#ma = nk.machine.FFNN(hilbert=hi, layers=layers)

ma.init_random_parameters(seed=1234, sigma=0.001)

# Sampler
sa = nk.sampler.MetropolisLocal(machine=ma)

# Optimizer
op = nk.optimizer.Sgd(1e-2)

# Quantum State Reconstruction
spvsd = nk.supervised.supervised(
    sampler=sa,
    optimizer=op,
    batch_size=64,
    niter_opt=10,
    output_file="output",
    samples=training_samples,
    targets=training_targets)

spvsd.run()

import numpy as np

mpi_rank = nk.MPI.rank()

# Load the data
N = 10
hi, rotations, training_samples, training_bases, ha, psi = generate(
    N, n_basis=2 * N, n_shots=500
)

# Machine
ma = nk.machine.RbmSpinPhase(hilbert=hi, alpha=1)
ma.init_random_parameters(seed=1234, sigma=0.01)

# Sampler
sa = nk.sampler.MetropolisLocal(machine=ma)

# Optimizer
op = nk.optimizer.AdaDelta()

# Quantum State Reconstruction
qst = nk.unsupervised.Qsr(
    sampler=sa,
    optimizer=op,
    batch_size=1000,
    n_samples=1000,
    rotations=rotations,
    samples=training_samples,
    bases=training_bases,
    method="Sr",
)

# Constructing a 1d lattice
g = nk.graph.Hypercube(length=20, n_dim=1)

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

# Hamiltonian
ha = nk.operator.Ising(h=1.0, hilbert=hi)

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

# Sampler
sa = nk.sampler.MetropolisLocal(machine=ma)
sa.seed(SEED)

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

# Variational Monte Carlo
vmc = nk.variational.Vmc(
    hamiltonian=ha,
    sampler=sa,
    optimizer=op,
    n_samples=1000,
    diag_shift=0.1,
    method="Sr",
)

vmc.run(output_prefix="test", n_iter=300, save_params_every=10)

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

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

# Ising spin hamiltonian at the critical point
ha = nk.operator.Ising(h=3.0, hilbert=hi)

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

# Metropolis Local Sampling
sa = nk.sampler.MetropolisLocal(machine=ma)

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

# Stochastic reconfiguration
gs = nk.variational.Vmc(
    hamiltonian=ha,
    sampler=sa,
    optimizer=op,
    n_samples=1000,
    diag_shift=0.1,
    method="Sr",
)

gs.run(output_prefix="test", n_iter=1000)