neuralqx.callbacks package

This module implements various VMC compatible callbacks.

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.

class LiveMonitoringCallback(host='127.0.0.1', port=8050, *, objective_label='Constraint')

Bases: object

NetKet callback that streams scalar metrics into a LiveMonitor.

This callback is designed to be registered with a NetKet/neuraLQX driver. On each iteration it extracts a scalar objective statistic and any scalar observables from the driver log, then enqueues them into the live monitor. If the UI requested an abort, it returns False to ask the driver to stop.

Parameters:

  • host (str) – Host/interface for the live monitor server.

  • port (int) – Preferred starting port for the live monitor server.

  • objective_label (str) – Display label for the objective/loss trace. On non-master MPI ranks, monitoring is disabled (monitor=None).

class GradientGradientCosineCallback(A, B, display=True)

Bases: Pytree

Compute the gradient-gradient cosine between two operators during optimisation.

At each driver step, the callback computes gradients of the expectations of A and B with respect to the variational parameters, flattens them, and evaluates

cos(phi) = Re( <gA, gB> ) / (||gA|| * ||gB|| + eps)

Interpretation guide - cos(phi) close to 1 suggests little or no frustration and gradients align - cos(phi) close to -1 suggests strong frustration and gradients oppose - cos(phi) close to 0 suggests near orthogonality

On the global MPI master process, the callback stores results in log_data under the keys “GradientGradientCosine” and “expectB”. When display is True, it also prints the cosine, both gradient norms, and the expectation value of B.

Parameters:

  • A (AbstractOperator) – First operator whose expectation gradient is used.

  • B (AbstractOperator) – Second operator whose expectation gradient is used.

  • display (bool) – If True, print a summary each step. If False, only log values.

Returns:

None.

Submodules