neuralqx.callbacks.pairwise_gradient_cosine_callback module

EXPERIMENTAL Callback for monitoring gradient alignment between two operators during VMC.

This module defines a NetKet driver callback that computes the cosine of the angle between the parameter gradients of two operators A and B at every optimisation step. It is useful for diagnosing frustration between objectives, for example between a constraint and an added penalty term.

The callback can both log values into the driver log data and optionally print a readable summary on the global MPI master process.

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.