diff --git a/example.py b/example.py new file mode 100644 index 0000000..d15341a --- /dev/null +++ b/example.py @@ -0,0 +1,106 @@ +# --- +# jupyter: +# jupytext: +# formats: ipynb,py:light +# text_representation: +# extension: .py +# format_name: light +# format_version: '1.5' +# jupytext_version: 1.15.1 +# kernelspec: +# display_name: Python 3 (ipykernel) +# language: python +# name: python3 +# --- + +# + +# this example can be opened as a notebook through jupytext +# - + +import netket as nk +import flax.linen as nn +import jax +import jax.numpy as jnp +from netket.utils.types import DType + +# + +# implements a Harmonic Oscillator with a single parameters + +# + +hi = nk.hilbert.Particle(1, D=1, pbc=False) + +ekin = nk.operator.KineticEnergy(hi, mass=1) + + +def vfun(x): + return 0.5 * jnp.sum(x**2) + + +epot = nk.operator.PotentialEnergy(hi, vfun) + +ham = ekin + epot +ham_t = lambda t: ham + + +# + +class HOModel(nn.Module): + param_dtype: DType = jnp.complex128 + + # See thesis of Giuseppe Carleo + # https://iris.sissa.it/retrieve/dd8a4bf7-04b2-20a0-e053-d805fe0a8cb0/1963_5357_carleo.pdf#page=94.12 + @nn.compact + def __call__(self, x): + x = x.reshape(-1) # single particle in 1d + a = self.param("alpha", jax.nn.initializers.ones, (), self.param_dtype) + return -a * x**2 / 2 + + +ma = HOModel() +ma +# - + +sa = nk.sampler.MetropolisGaussian(hi, sigma=1.0, sweep_size=hi.size * 2, n_chains=128) +vs = nk.vqs.MCState(sa, ma, n_samples=8 * 1024, n_discard_per_chain=4) +vs + +vs.samples.shape + +vs.sampler_state.acceptance + +from tre_tvmc.driver import QDynamics +from tre_tvmc.tre import TREGenerator +from tre_tvmc.solver import block_solver +from jax.scipy.sparse.linalg import cg + +# + +tre_generator = TREGenerator(hi, ham, order=4) +dt = 1e-2 +Tmax = 1 + +solver = block_solver(cg, vs.parameters) +qgt = nk.optimizer.qgt.QGTJacobianPyTree +sr = nk.optimizer.SR(qgt, holomorphic=True, diag_shift=1e-3) +opt = nk.optimizer.Sgd(1e-2) + +qgt +# - + +method_kwargs = { + "sample_Upsi": False, + "n_sigma_check": 0.1, + "n_redo": 20, + "cv_coeff": -0.5, +} +te = QDynamics( + ham_t, + vs, + dt, + tevo_generator=tre_generator, + preconditioner=sr, + optimizer=opt, + method_kwargs=method_kwargs, +) + +te.run( + Tmax, n_iter=25, out="runs/test", out_inner="runs/test/", show_inner_progress=False +)