How to use MetropolisSamplerNumpy

[whatisspin]

Hello,

I was trying to make a custom sampler by defining the transition rule. While it works with MetropolisSampler, it does not with MetropolisSamplerNumpy. For example,

import netket as nk
from netket.sampler import MetropolisRule
from netket.utils.struct import dataclass

@dataclass
class sampler_custom(MetropolisRule):
    def transition(rule, sampler, machine, parameters, state, key, sigma_chain):

        return sigma_chain,None

hi= nk.hilbert.Spin(0.5,10)
ma= nk.models.RBM()
sampler1= nk.sampler.MetropolisSampler(hilbert=hi, rule= sampler_custom(), n_chains=5)
sampler2= nk.sampler.MetropolisSamplerNumpy(hilbert=hi, rule= sampler_custom(), n_chains=5)

vs1= nk.vqs.MCState(sampler1, ma, n_samples=10)
vs2= nk.vqs.MCState(sampler2, ma, n_samples=10)

print(vs1.samples)
print(vs2.samples)

gives output,

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

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

 [[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]]

What am I doing wrong?

Thank You.

[PhilipVinc]

The documentation of MetropolisSamplerNumpy says that this sampler is executed through Numpy instead of Jax that is the standard throught netket.

Therefore, it requires the transition kernels to be written in NumPy instead of Jax. Admittedly this is not mentioned anywhere, and we could improve the documentation. To write a NumPy Kernel, start from one of the existing numpy kernels for inspiration such as LocalRuleNumpy or HamiltonianNumpy or the very generic CustomNumpy.

The reason for the existence of this numpy sampler is that when we started the rewrite of NetKet 3 in Jax, we could not write some transition rules in Jax efficiently (in particular HamiltonianRule and Custom) and therefore we attempted to have this different implementation which could sidestep the limitation.

As a lot happened since then, we should probably revisit whether we still need to have this dumpy-specific implementation or if this can be dropped.
In general, however, I'd expect worse performance from this one except in some very specific use cases.

May I ask what transition rule you are trying to implement?

[PhilipVinc]

To be more precise, if you look at how numpy transition rules are implemented they modify in-place the fields state.sigma1 and state.log_prob_corr and they do not return anything.

This is necessary in order for the samplers to be efficient.

So a correct implementation of your rule would be...

@dataclass
class sampler_custom(MetropolisRule):
    def transition(rule, sampler, machine, parameters, state, key, sigma_chain):
       np.copyto(state.σ1, sigma_chain)

[whatisspin]

Thank you for the very quick answer.

I have a custom Hilbert space with some symmetry related constraints and I want my transition rule to respect those constraints. I could not do this without using conditional statements. (Maybe I can create a minimum example and post here later.)

I could implement that in Jax using its control flow primitives, which was not fun. The reason I am trying to go to NumPy is because for larger systems (>20 sites), the Jax sampler is very slow. For 40 sites or more, it takes hours for the optimization to start. I am still not sure what is the reason for this and wanted to compare it with a NumPy sampler. (I have narrowed it down to jax.random.split I am using inside the body function of a jax.lax.fori_loop. I do not understand it enough to say more about it.)

[PhilipVinc]

Indeed, you should absolutely not use jax control flow (cons) for this thing, because of the way it is implemented.
Try replacing the jax.lax.cond with jnp.where, it will surely run much faster, especially on GPU, assuming your logic is not extraordinarily convoluted.

In general, the best way to implement this things in jax is to try to write them as pure linear algebra operations. A condition can always be expressed condition*value_true + jnp.logical_not(condition)*value_False ((which is what jnp.where does, to some extent).

Also, yes, don't use fori_loop in the sampler if you can...

From what you tell, indeed, you'd be better off using numpy sampler.

[PhilipVinc]

still I'd be curious to see what you come up with.