Starting sampler with a specific state?

[fabienalet]

Hi all,
I have a basic question on how to start a sampler with a specific state, say the Neel state on four sites
neel = [1,0,1,0]
for a given machine ma and sampler sa
I tried
sampler.reset(machine=ma,state=neel)
but I'm missing the correct argument parameters (I tried several things but it failed).
Any simple workaround ? Thanks !

[PhilipVinc]

State is not a ket, but rather a structure holding the state of the sampler (which also includes the current batch of kets) We could add what you say though, it is indeed a valid feature. A non Hacky workaround is to define a new Rule (like local, or exchange) inheriting from whatever rule you want, and override the function `random_state` to give you a batch of states you want. I’ll give you an example in a bit. Il 3 set 2021, 17:36 +0200, Fabien Alet ***@***.***>, ha scritto:

…

Hi all, I have a basic question on how to start a sampler with a specific state, say the Neel state on four sites neel = [1,0,1,0] for a given machine ma and sampler sa I tried sampler.reset(machine=ma,state=neel) but I'm missing the correct argument parameters (I tried several things but it failed). Any simple workaround ? Thanks ! — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub, or unsubscribe.

[PhilipVinc]

here.

Just consider that metropolis sampler resets the state only once by default, at the beginning of the chain (reset_chain=False by default).
If you want to re-reset it every time you change the parameters (so at every VMC step) you should pass
reset_chain=True to MetropolisSampler.

You can derive from any rule you want to use.

import netket as nk
import jax.numpy as jnp
from netket.utils import struct

# 1D Lattice
L = 20  # 10

g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)

# Hilbert space of spins on the graph
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)

# Ising spin hamiltonian
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)

# RBM Spin Machine
ma = nk.models.RBM(alpha=1, use_visible_bias=True, dtype=float)

# Special rule
@struct.dataclass
class SpecialLocalRule(nk.sampler.rules.LocalRule):

    def random_state(
        rule,
        sampler,
        machine,
        parameters,
        state,
        key,
    ):
        """
        Generates a random state compatible with this rule.

        By default this calls :func:`netket.hilbert.random.random_state`.

        Arguments:
            sampler: the sampler
            machine: the function to evaluate the model
            parameters: the parameters of the model
            state: the current sampler state
            key: the PRNGKey to use to generate the random state
        """
        
        # Default
        # return sampler.hilbert.random_state(
        #     key, size=sampler.n_batches, dtype=sampler.dtype
        # )

        # only works on even spaces. Otherwise find another way to do it...
        # or hardcode it.
        assert sampler.hilbert.size % 2 == 0
        neel_state = jnp.array([0,1]*(sampler.hilbert.size//2))

        return jnp.tile(neel_state, (sampler.n_chains_per_rank, 1))

    def __repr__(self):
        return "SpecialMagicLOcalRule"


# Metropolis Local Sampling
sa = nk.sampler.MetropolisSampler(hi, SpecialLocalRule(), n_chains=16, reset_chains=False)

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.1)

# Variational monte carlo driver
gs = nk.VMC(ha, op, sa, ma, n_samples=1000)

# Run the optimization for 300 iterations
gs.run(n_iter=300, out="test")

[fabienalet]

Hi @PhilipVinc
Thanks a lot, this helped me solving my issue (and understanding better the structure). As I used a Hamiltonian sampler, I had to adapt things a bit and build a different sampler (which I show here for the record)

from netket.sampler.metropolis import MetropolisSampler

def NeelMetropolisHamiltonian(hilbert, hamiltonian, *args, **kwargs) -> MetropolisSampler:

      rule = netket.sampler.rules.HamiltonianRule(hamiltonian)

      def random_state(rule,
        sampler,
        machine,
        parameters,
        state,
        key,
        ):
           assert sampler.hilbert.size % 2 == 0
           neel_state = jnp.array([0,1]*(sampler.hilbert.size//2))
           return jnp.tile(neel_state, (sampler.n_chains_per_rank, 1))

      return MetropolisSampler(hilbert, rule, *args, **kwargs)

Thanks again!

[PhilipVinc]

You're welcome.

Though I think that you forgot something in the snippet above?
You define random_state but it's not associated with a class or anything.
For posterity, it would be nice if you could edit your message above so that it actually works?