Understanding the states in constrained Hilbert spaces

[waleed-sh]

I am currently working with a spin-like system and a Hilbert space which is constrained such that it only includes states where the basis label of two specific lattice sites sum up to zero. For example, a constraint function which looks like

def conFunc(states):
    return [state[0] == -state[3] for state in states]

I was able to implement a Metropolis Local sampler which respects this constraint. Essentially, it does everything the Metropolis Local sampler does, but then before returning the proposed configurations it manually sets the values of these specific lattice sites to the same value but opposite sign (see below).

My questions are:

1. how do I understand the variational state now? Would I be really working only in a subspace of the entire Hilbert space for example when taking expectation values, etc.?
2. using the .states_to_numbers() method with a state which does not satisfy this condition still returns a a number for it. This does not make sense, I thought that the space now does not contain these states. However, using .numbers_to_states() does not allow me to go over the actual number of states in this constrained space.

Edit: Was able to test whether the sampler works correctly or not, but still struggling to understand the states... Any help would be appreciated!

One additional question regarding the sampler (related to this #1759 ): Is the manner in which the sampler is implemented now correct in the sense that it will generate only states satisfying the criteria I specify? Essentially, I just added the last line after flipping the state

σp, _ = nk.hilbert.random.flip_state(hilb, key_flip, σ, indxs)
σp = self.reconfigure(σp)

to the transition() method of the sampler, and then defined

def reconfigure(self, σp):
    def update_σi(σi):
        σi = σi.at[3].set(-σi[0])
        return σi

    return jax.vmap(update_σi)(σp)

I also updated the random_state() method accordingly. I think the proposed configurations are all "reconfigured" to satisfy this criteria and the sampler always proposes states which are reconfigured in this manner, or am I wrong? Would this then be an alternative approach to #1759?

[PhilipVinc]

For example, a constraint function which looks like

This seems a valid choice.

how do I understand the variational state now? Would I be really working only in a subspace of the entire Hilbert space for example when taking expectation values, etc.?

Yes. The idea is that your network can take over any value outside the subspace of interest, because we will only be evaluating it inside of the subspace when evaluating expectation values.

using the .states_to_numbers() method with a state which does not satisfy this condition still returns a a number for it.

I think this is a bug. It will be fixed in the upcoming version of netket, and will give an hard error.

One additional question regarding the sampler (related to this https://github.com/orgs/netket/discussions/1759 ): Is the manner in which the sampler is implemented now correct in the sense that it will generate only states satisfying the criteria I specify? Essentially, I just added the last line after flipping the state

This seems correct as well.

[waleed-sh]

Thank you for the reply! It helped clarify things.

[PhilipVinc]

By the way, yesterday evening we released a new version of netket that changed a bit how we deal with the constraints internally.

From your point of view, what is now necessary is that the constraint must be a function that can be jitted with jax.
We will add more documentation soon, but in case you need some help, do open an issue.

In particular, for

def conFunc(states):
    return [state[0] == -state[3] for state in states]

this will not work, and you should rather write it as something like

def conFunc(states):
 return states[...,0] == -states[...,3]

if you really cannot write it in jax, you can use a callback like

def conFuncPy(states):
     # cannot express this in jax
     return ...

def conFunc(states):
    return jax.pure_callback(conFuncPy, jax.ShapeDtypeStruct(states.shape[:-1], bool), states)

[PhilipVinc]

It will also, now, throw an error in the cases where it should.

[waleed-sh]

Thanks for the heads-up! The constraint functions I use are so far as simple as the one I wrote so it's fine but will keep that in mind!