Define and document an interface for publicly defining custom constraint

[PhilipVinc]

As some people have been asking this over time, and we have finally a much cleaner interface thanks to the work of @inailuig , I feel safer formalising this partially.

In particular, this PR:

• defines the interface that discrete Hilbert constraint must follow
• Adds a keyword argument to Spin and Fock to define a custom constraint
• Formalises how to define random_state for general constraints (using dispatch rules)

If this works well, in the future we might start throwing warnings when we have some combinations of Hilbert space/constraints/samplers that we know do not work together...

[PhilipVinc]

This PR is now ready. It introduces two new interfaces

1. a way to specify an arbitrary constraint for the existing Hilbert Spaces by doing

nk.hilbert.Spin(0.5, N=10, constraint=MyConstraint(...))

2. A way to specify how to generate random samples for arbitrary constraints, by specifying

@nk.hilbert.random.random_state.dispatch
def random_state(hi: Spin, constraint: MyConstraint, key, shape :int, dtype):
 ...

Those two things together will allow us to make it finally simple to support arbitrary different constraints without forcing the users to define entire new hilbert spaces, which is more involved.
This also promotes our constraint types to full-fledged public types, which we could start using to throw warnings when a constraint is not compatible with a sampler...

What do you think?

[gcarleo]

looks quite neat! is it limited to Spin and Fock as originally done or is it more general? In any case, I have a likely trivial doubt concerning having the constraint as second argument
hi: Spin, constraint: MyConstraint , is it positional ?

[gcarleo]

also, maybe we should remove "Defining Custom constraints" from the Doc and only keep the new section?

[PhilipVinc]

is it limited to Spin and Fock as originally done or is it more general?

It's for all HomogeneousHilbert spaces.

In principle we could scale it up to composite systems as well in the future, but would require some more work on TensorHilbert.

In any case, I have a likely trivial doubt concerning having the constraint as second argument

Yes it's a positional argument. This is just a function the user has to define, but will never use it directly. The user will simply do

my_hilbert = nk.hilbert.Spin(0.5, 5, constraint=MyConstraint(...))
my_hilbert.random_state(jax.random.key(1), 10)

If the user forgets to define the custom dispatch he will get this very informative error

         You are using a custom constraint. You must define how to generate random states
         for this particular state.
         To do this, define
  
       @nk.hilbert.random.random_state.dispatch
         def random_state(hilb: nk.hilbert.Spin, constraint: MyConstraint , key, batches: int, *, dtype=None):
             return ...
  
         and the function should return an array of size `(batches, hilb.size)` and dtype `dtype` (if dtype is none you can pick anything you think reasonable). The function will be `jax.jit`ted, so if you cannot write it using only jax functions you should use `jax.pure_callback` to use standard python functions.

Which I think is informative enough