Getting latest samples from VMC training

Created by: VolodyaCO

Short question

How do I get the latest configurations of the training phase of a Vmc gs = nk.Vmc(hamiltonian=ha, sampler=sa, optimizer=op, n_samples=n_samples, sr=sr)? gs._samples seems to be the correct variable, but it is not.

Long question

I am trying to solve a Hamiltonian of 5 completely connected spins. The Hamiltonian has a 5-fold degenerate ground state, where any of the spins is up, and the rest is down. The implementation of the Hamiltonian is shown next:

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

L = 5
g = nx.complete_graph(L)
g = nk.graph.NetworkX(g)

hi = nk.hilbert.Spin(s=0.5, graph=g)

sz = np.array([[1, 0], [0, -1]])
ha = nk.operator.LocalOperator(hi)
for i in range(L):
    ha += nk.operator.LocalOperator(hi, 0.5*(L-2)*sz, [i])
    for j in range(L):
        if i != j:
            ha += nk.operator.LocalOperator(hi, 0.25*np.kron(sz, sz), [i, j])

ma = nk.machine.RbmSpin(alpha=1, hilbert=hi)
sa = nk.sampler.MetropolisLocal(ma, n_chains=32, sweep_size=1)
#sa = MetropolisLocalExchange(ma, n_chains=32, sweep_size=None)
op = nk.optimizer.Sgd(ma, learning_rate=0.05)
sr = nk.optimizer.SR(ma, diag_shift=0.1)
gs = nk.Vmc(hamiltonian=ha, sampler=sa, optimizer=op, n_samples=100, sr=sr)
gs.run(n_iter=250, out="test")

After I run this, the energy converges correctly to the -3.5 value that it should. Now, I want to see which spin configurations are relevant to the ground state. For this, I check which entries of the complete machine state are large:

from qutip import state_index_number
psi_ma = ma.to_array()

idx = np.where(np.abs(psi_ma)**2>=1e-5)[0]
for i in idx:
    print(state_index_number([2]*L, i), "with coeff", (np.abs(psi_ma)**2)[i])

I got this:

[0, 0, 0, 1, 1] with coeff 1.4798755296251725e-05
[0, 1, 0, 0, 1] with coeff 1.4001762026219615e-05
[0, 1, 0, 1, 0] with coeff 1.5824190892762988e-05
[0, 1, 0, 1, 1] with coeff 0.004003922485146549
[0, 1, 1, 1, 1] with coeff 1.6373536077475163e-05
[1, 0, 0, 0, 1] with coeff 1.2678379200803993e-05
[1, 0, 0, 1, 0] with coeff 1.4328560384662225e-05
[1, 0, 0, 1, 1] with coeff 0.003625489953496893
[1, 0, 1, 1, 1] with coeff 1.4825983962557237e-05
[1, 1, 0, 0, 0] with coeff 1.3556889661874039e-05
[1, 1, 0, 0, 1] with coeff 0.003430237647768317
[1, 1, 0, 1, 0] with coeff 0.003876707463259443
[1, 1, 0, 1, 1] with coeff 0.9809055190037385
[1, 1, 1, 0, 1] with coeff 1.4027524281103155e-05
[1, 1, 1, 1, 0] with coeff 1.585330628826408e-05
[1, 1, 1, 1, 1] with coeff 0.004011289420208237

These configurations are read as 0 -> up spin, 1 -> down spin. Clearly, there is one state [1, 1, 0, 1, 1] with coeff 0.9809055190037385 that gets almost all the probability. This is the correct solution (there are 5 correct solutions). However, this makes that all the other states get "noisy probabilities". Note that the state [0, 1, 1, 1, 1] with coeff 1.6373536077475163e-05 is also a solution, but it was given a small probability (because the algorithm converged first to the other configuration. This is not a problem at all for what I need. However, I have extracted the probabilities by sampling the complete Hilbert space with ma.to_array(). I want, instead, to sample from netket's sampler to get a bunch of configurations, among which, I hope, there is a solution to the problem I'm looking for. Since the energy converged correctly, I assume that the latest samples contain the configuration [1, 1, 0, 1, 1].

The question is, then, how do I retrieve this last batch of samples? For instance, gs._samples contains only these configurations (I print the unique ones):

[ 1.  1. -1. -1.  1.]
[ 1. -1. -1.  1.  1.]
[-1.  1. -1.  1.  1.]
[ 1.  1. -1.  1.  1.]

(Sorry for the change in notation, when 1 and -1 appear, those are spins. when 0 and 1 appear, 0 -> spin +1, 1 -> spin -1)