neuralqx.experimental.driver package

Experimental VMC drivers in neuraLQX.

class MultiStateVMC(variational_state, hamiltonian, optimizer, *, preconditioner=<netket.optimizer.preconditioner.IdentityPreconditioner object>, lambda_ortho=1.0)

Bases: VMC

VMC driver for MultiMCState with an optional orthogonality penalty.

For each contained state, this driver computes energy statistics and gradients of the Hamiltonian objective. If lambda_ortho is nonzero and more than one state is present, it additionally applies a pairwise penalty based on a fidelity-like overlap estimator to encourage the states to become mutually orthogonal.

Preconditioning is applied independently per state. Interpreted as a joint optimisation problem over all parameters, this corresponds to a block-diagonal preconditioner.

Parameters:

  • variational_state (MultiMCState) – Multi-state variational object containing multiple independent states.

  • hamiltonian (Union[AbstractOperator, list]) – Operator or list of operators defining the objective for each state.

  • optimizer (Any) – Optimiser used to update parameters from the (preconditioned) gradients.

  • preconditioner (Callable[[VariationalState, Any, Any | None], Any]) – Preconditioner applied per state to transform raw gradients.

  • lambda_ortho (float) – Strength of the orthogonality penalty. Set to 0 to disable.

Raises:

TypeError – If any operator acts on a Hilbert space incompatible with the variational state.

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