Add block-spin update sampler

[kevinchern]

This PR introduces a block-spin update sampler.
It includes duplicate code from #40

[jackraymond]

Great job, I'll come back to this when you add tests. Per your summary. Thanks for getting this draft up quickly.

Would be nice to have a tidy public GPU code to use and cite for speedup relative to our CPU implementation in higher throughput applications as soon as possible.

[kevinchern]

Samplers should have a consistent interface. Related: #52, #47

[kevinchern]

@jackraymond I updated the PR:

1. added tests (some notable mentions below)
   • test gibbs update by checking magnetization is close to expected
   • test metropolis update by checking magnetization is close to expected
   • test metropolis update by checking oscillatory behavior for beta = 0 case
2. added documentation (first pass)
3. capitalized Metropolis/Gibbs and also make it case-insensitive
4. The sampler has a simplified signature. It no longer tries to do the colouring itself. Instead, it requests for a colouring function. This generalizes it beyond CPZ topologies. One could, for example, default to a colouring heuristic when no colouring is given, but I've opted away from that for modularity.
   5. Now a metadata dictionary is returned along with the sample.

[VolodyaCO]

I left some comments about form.

Overall looks fantastic. Thank you.

[VolodyaCO]

Can the user compile a method of an existing nn.Module subclassed instance? Ideally, we would want the user to decide whether they want to compile the _step method. It seems as though torch.compile(module) compiles the __call__ method, so maybe we can wrap _step in it and time the compiled and non-compiled versions.

[kevinchern]

Addressed @thisac and @VolodyaCO 's PRs. Briefly,

1. Remove inheritance of nn.Module
2. Add functional module docstring
3. Add explicit named args for, and renamed to randspin
4. Improved documentation and type-hinting

[jackraymond]

Looks good tests ran for me.

I'd recommend you allow setting of x as an initial condition. Check it doesn't move when num_sweeps is none, and check it oscillates under BlockMetropolis when temperature is infinity as tests perhaps.

[kevinchern]

Looks good tests ran for me.

I'd recommend you allow setting of x as an initial condition. Check it doesn't move when num_sweeps is none, and check it oscillates under BlockMetropolis when temperature is infinity as tests perhaps.

In the latest commit, I made the following changes:

• Fixed documentation RE nonzero & finite temperature regimes.
• Allow setting of initial states.
  • Test by checking sampler.x.tolist() is equal to input initial states list
• Added custom errors to distinguish from ValueError

@jackraymond @thisac

[VolodyaCO]

All comments have been addressed. Thank you @kevinchern

[thisac]

Thanks @kevinchern! 🎄

[thisac]

Some cleanup needed. Spaces around - and *, and perhaps separating out the torch.randint() call since the final 2* and -1 get kind-of lost in there.

[thisac]

Also not convinced that this i needed. Couldn't you just use the regular torch.randint() and pass the results through bit2spin_soft()? Other than that this only sets low/high which isn't too cumbersome to do manually.

[kevinchern]

-1 get lost

Changed to b = randint ... and return 2 * b - 1

[kevinchern]

Also not convinced that this i needed. Couldn't you just use the regular torch.randint() and pass the results through bit2spin_soft()? Other than that this only sets low/high which isn't too cumbersome to do manually.

My motivation for adding it is similar to adding a bit2spin function. It's very simple to write repeatedly, but it helps with readability to have a dedicated function. What are some reasons to avoid adding this?

[thisac]

Alternatively, either be clear about this being a special case of torch.randint(), and simplify it a bit to make it clear that it should work with all kwargs supported in randint.

def randspin(size: _size, **kwargs) -> torch.Tensor:
    """Wrapper for ``torch.randint()`` restricted to spin outputs."""
    if kwargs.get("high", None) or kwargs.get("low", None):
        raise ValueError(...)

    b = torch.randint(0, 2, size, **kwargs)
    return 2 * b - 1

or remove the kwargs from the signature and add all the relevant arguments, which circumvents the need to check for "high" and "low".

[thisac]

Better put in a setUpClass() method.

[kevinchern]

If I understood correctly, that would be incompatible with @parameterized.expand

[thisac]

Generally, you should try to avoid calling hidden methods/attributes even in tests (only what is meant to be public should need to be tested), but need to think about this whether it makes sense here.

Overall, not sure if this test could be streamlined a bit.