Automatic checkpointing

[PhilipVinc]

A long way to go...

Will need some help to make this play nice with sharding and MPI.

Must decide what to do with loggers. Right now we override the old ones. Can't really concatenate.

temporary interface.

from etils import epath
import orbax.checkpoint as ocp
import numpy as np

# Keeps a maximum of 3 checkpoints, and only saves every other step.
path = epath.Path('/tmp/nkcheck3')
checkpointer = ocp.CheckpointManager(
    path,
    options=ocp.CheckpointManagerOptions(max_to_keep=3, save_interval_steps=1),
)

import netket as nk
import optax

L = 10
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)
ma = nk.models.RBM(alpha=1, param_dtype=float)
sa = nk.sampler.MetropolisLocal(hi, n_chains=16)
op = nk.optimizer.Sgd(learning_rate=optax.linear_schedule(0.1, 0.0001, 500))
sr = nk.optimizer.SR(diag_shift=0.01)
vs = nk.vqs.MCState(sa, ma, n_samples=1008, n_discard_per_chain=10)

# Variational monte carlo driver with a variational state
gs = nk.VMC(ha, op, variational_state=vs, preconditioner=sr, checkpointer= checkpointer)

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

if you run it twice, the second run goes from 50 to 100 steps

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

Attention: Patch coverage is 34.78261% with 30 lines in your changes are missing coverage. Please review.

Project coverage is 81.45%. Comparing base (9b069f1) to head (d72efb2).

Files	Patch %	Lines
netket/driver/abstract_variational_driver.py	27.77%	24 Missing and 2 partials ⚠️
netket/experimental/driver/tdvp_common.py	20.00%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1667      +/-   ##
==========================================
- Coverage   82.98%   81.45%   -1.54%     
==========================================
  Files         303      303              
  Lines       18433    18476      +43     
  Branches     2718     2726       +8     
==========================================
- Hits        15296    15049     -247     
- Misses       2465     2791     +326     
+ Partials      672      636      -36     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.