neuralqx.experimental.driver package

Experimental VMC drivers in neuraLQX.

class SingleTrunkMultiHeadVMC(variational_state, hamiltonian, optimizer, *, preconditioner=<netket.optimizer.preconditioner.IdentityPreconditioner object>, lambda_ortho=1.0, energy_weights=None, enforce_machine_pow_2=True, preconditioner_state_index=0)

Bases: VMC

VMC driver for a shared-parameter ST-MH ensemble built from head-specific MCState views.

The optimized objective is

\[C(\Theta) = \sum_k w_k E_k(\Theta) + \lambda \sum_{i<j} F_{ij}(\Theta),\]

where F_ij is the normalized pairwise fidelity estimator (for machine_pow=2). The sum_{i<j} form is equivalent to (1/2) sum_{i!=j} used in the paper.

The existing MultiStateVMC is correct for MT-MH (independent parameter sets) and preconditions each state separately. This driver is for ST-MH. It aggregates all energy and penalty gradients into one shared parameter pytree and applies the optimizer/preconditioner exactly once.

Note

The default identity_preconditioner is exact for Euclidean gradient descent on the chosen Monte Carlo objective. If you pass an SR/QGT preconditioner, the driver applies it using a single reference head state’s geometry (preconditioner_state_index), which is a practical approximation.

Submodules