Sampling from MCMixedState

[VolodyaCO]

I trained a model that looks like this:

>>> vstate: nk.vqs.MCMixedState = load_best_state(3)
>>> print(vstate)
MCMixedState(hilbert = DoubledHilbert(Fock(n_max=4, N=3)), sampler = MetropolisSampler(rule = LocalRule(), n_chains = 16, n_sweeps = 6, reset_chains = False, machine_power = 2, dtype = <class 'numpy.float64'>), n_samples = 608)

Here, load_best_state is just a function that loads an .mpack file from a previous experiment onto the state vstate.
As you can see, the problem is defined in the Fock space of three boson modes. I want to see the mean occupancy of each mode:

>>> x = vstate.diagonal.samples.reshape(-1, 3)
>>> print("Mean occupancy of each Fock mode:", np.mean(x, axis=0))
Mean occupancy of each Fock mode: [1.4875     1.46666667 1.19583333]

Now, I want to continue sampling in order to, say, get a better approximation of the mean occupancy, but I see that the diagonal samples-which are the only ones used if I want to do vstate.expect-drift towards 0:

>>> # Create new sample:
>>> x = vstate.diagonal.sample().reshape(-1, 3)
>>> print("Mean occupancy of each Fock mode:", np.mean(x, axis=0))
Mean occupancy of each Fock mode: [1.26666667 1.15       1.125     ]
>>> # Repeat the process a couple more times:
>>> x = vstate.diagonal.sample().reshape(-1, 3)
>>> print("Mean occupancy of each Fock mode:", np.mean(x, axis=0))
Mean occupancy of each Fock mode: [0.57083333 0.50416667 0.49166667]
>>> x = vstate.diagonal.sample().reshape(-1, 3)
>>> print("Mean occupancy of each Fock mode:", np.mean(x, axis=0))
Mean occupancy of each Fock mode: [0.55833333 0.48333333 0.39166667]

Why is this? Is this to be expected? I would expect the samples to be very similar to the last batch of samples during training.

[PhilipVinc]

Most likely the chains are not at convergence. If I’m right, you’d notice by using an operator to compute those expectation values and you’d see that the split Rhat is off.

This might happen if you are loading only the parameters but not the whole variational state (and it’s sampler state) , because then you’d be resetting the Markova chain and it might take a while to thermalise.

try running a long chain and discard it before looking at this

[VolodyaCO]

I'm not entirely sure about this. Just after loading the state from the .mpack file, I'm able to compute expected values:

vstate: nk.vqs.MCMixedState = load_best_state(3)
print(vstate.diagonal.samples.shape)
for num_op in photon_number_observable(vstate.hilbert_physical):
    print(num_op)
    print(vstate.expect(num_op))

yields

(15, 16, 3)
LocalOperator(dim=3, acting_on=[(0,)], constant=0.0, dtype=float64)
1.488+0.000j ± 0.096 [σ²=2.233, R̂=1.3872]
LocalOperator(dim=3, acting_on=[(1,)], constant=0.0, dtype=float64)
1.467+0.000j ± 0.095 [σ²=2.182, R̂=1.3342]
LocalOperator(dim=3, acting_on=[(2,)], constant=0.0, dtype=float64)
1.196+0.000j ± 0.091 [σ²=1.966, R̂=1.3141]

I see that Rhat is in all cases far from 1, but it looks like there's already a sample in the sampler, isn't it?

[PhilipVinc]

when R_hat> 1.02 (and especially if it's as high as 1.3) the MCMC is not converged.
In that case, you should in theory disregard everything else. The expectation value, the error estimate, the variance.... everything!
You should not look at the numbers. They are wrong. The algorithms used to estimate them are wrong if Rhat tells you that you are not converged.

So the fact that your exp value is changing is.. not important in that case. You should concentrate on why your MCMC is not at convergence.

[VolodyaCO]

Ok yes this makes sense. Thanks!