Update README.md

README.md

@@ -18,6 +18,61 @@ The computation has been distributed over 4 A100-64GB GPUs for few hours.
 Use the code below to get started with the model. In particular, we sample the model for a fixed value of the externa field \\(h\\) using NetKet.
 
 ```python
+from functools import partial
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+import netket as nk
+
+import flax
+from flax.training import checkpoints
+
+flax.config.update('flax_use_orbax_checkpointing', False)
+
+lattice = nk.graph.Hypercube(length=100, n_dim=1, pbc=True)
+
+revision = "main"
+h = 1.0 #* fix the value of the external field
+
+from transformers import FlaxAutoModel
+wf = FlaxAutoModel.from_pretrained("nqs-models/ising_fnqs", trust_remote_code=True)
+N_params = nk.jax.tree_size(wf.params)
+print('Number of parameters = ', N_params, flush=True)
+
+hilbert = nk.hilbert.Spin(s=1/2, N=lattice.n_nodes)
+hamiltonian = nk.operator.IsingJax(hilbert=hilbert, graph=lattice, h=h, J=-1.0)
+
+action = nk.sampler.rules.LocalRule()
+sampler = nk.sampler.MetropolisSampler(hilbert=hilbert, 
+                                       rule=action, 
+                                       n_chains=12000, 
+                                       n_sweeps=lattice.n_nodes)
+
+key = jax.random.PRNGKey(0)
+key, subkey = jax.random.split(key, 2)
+vstate = nk.vqs.MCState(sampler=sampler, 
+                        apply_fun=partial(wf.__call__, coups=h), 
+                        sampler_seed=subkey,
+                        n_samples=12000, 
+                        n_discard_per_chain=0,
+                        variables=wf.params,
+                        chunk_size=12000)
+
+from huggingface_hub import hf_hub_download
+path = hf_hub_download(repo_id="nqs-models/ising_fnqs", filename="spins", revision=revision)
+samples = checkpoints.restore_checkpoint(path, prefix="spins", target=None)
+samples = jnp.array(samples, dtype='int8')
+vstate.sampler_state = vstate.sampler_state.replace(σ = samples)
+
+import time
+# Sample the model
+for _ in range(10):
+    start = time.time()
+    E = vstate.expect(hamiltonian)
+    vstate.sample()
+
+    print("Mean: ", E.mean.real / lattice.n_nodes, "\t time=", time.time()-start)
 
 ```