neuralqx.callbacks.converged_state_export module

This module defines a NetKet-compatible callback for Variational Monte Carlo (VMC) optimisation runs. The callback monitors a chosen loss statistic (e.g. mean or error estimate) and triggers an export of the current Monte Carlo state whenever the monitored value reaches (or falls below) a user-defined target threshold.

The exported state is intended to be a fully serialised MCState (including the sampler state), written via the solver’s export facility (typically to .mpack format). This enables checkpointing “good” optimisation states at the moment a desired loss level is achieved.

The callback is implemented as a netket.utils.struct.Pytree, making it compatible with NetKet’s driver callback mechanism and JAX/pytrees where appropriate.

Notes

  • The callback always returns True to indicate that optimisation should continue, it is a logging/export side-effect hook rather than a stopping criterion.

  • Only a fixed set of monitoring keys is accepted to avoid silent typos and to align with the statistics provided by NetKet log entries.

class StateExportCallback(target, solver, monitor='Mean')

Bases: Pytree

Export the serialised MCState when the monitored loss statistic reaches a target threshold.

The callback reads the current loss entry from log_data[driver._loss_name], extracts the requested statistic (configured via monitor), and compares it to target. If the value is less than or equal to target, it exports the driver’s current state through the provided Solver instance. The export marker includes step and diagnostic statistics.

Parameters:

  • target (float) – Target threshold for the monitored loss statistic. When the monitored value is <= this threshold, an export is triggered.

  • solver (AbstractSolver) – Instance of neuralqx.solver.Solver used to perform the serialised export (must provide export_serialized_state(…)).

  • monitor (str) – Name of the statistic to read from the driver’s loss log entry. Must be one of [“Mean”, “Sigma”, “R_hat”, “TauCorr”]. Defaults to “Mean”.

Raises:

IncorrectMonitoringValueError – If monitor is not one of the allowed monitoring keys.