[ENH] Add function to check if state is in Hilbert space

[PhilipVinc]

For Renyi entropy calculation and a few other things it might be useful to have an utility function to check if a configuration )vector) is within the computational space.

This could be implemented very neatly by leveraging __contains__ to use it as

hi = nk.hilbert.Spin(0.5, 10)
xs = hi.all_states() # 2^N x N 

valid = xs in hi  # bool vector 2^N

or a more traditional

valid = hi.contains(xs)

class MyHilbert(DiscreteHilbert):
   def __contains__(self, x):
        valid_numbers = np.all(x > min(self.local_states) and x< max(self.local_states), axis=-1)
        if self.is_constrained:
             valid_numbers = np.logical_and(valid_numbers, self._constraint_fun(x))
        return valid_nmbers

The main question is: we can check three things

• (1) constraint respected if it is there
• (2) numbers in the interval of valid values (ie for spin 1/2, between -1 and 1) but not check that it is exactly one of those values
• (3) do the exact check, meaning that the value is a correct integer

Of course, doing the full check is more expensive and often we don't need to check it, as it's kind-of-guaranteed by our sampling.
However, sometimes we might want to do the full check

Possibly, we could have a keyword argument Hilbert.contains(x, *, fast=True/False)which defaults to False, in which case it does not check (3) but only (1) or (2), or even only (1)...

What makes sense to do?

Ps: once agreed upon, it would be neat to have a PR by someone!

[jwnys]

Related to this: in my opinion the MCState should initialize from a hilbert space sample, not a set of zeros that is often outside of the hilbert space (this is especially very annoying with fermions):

netket/netket/vqs/mc/mc_state/state.py

Line 282 in 9286307

dummy_input = jnp.zeros((1, self.hilbert.size), dtype=dtype)

[gcarleo]

Ah, indeed this is a bug then?

[PhilipVinc]

is it a thing in practice? I thought flax does not really evaluate the model at init... https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init

If you have a non flax model, however, this is an issue.

Indeed, just taking a random state would be a better idea..