neuralqx.experimental.driver.stmh_multi_vmc

fidelity_expect_and_grad_joint_tied(apply_fun_i, apply_fun_j, machine_pow_i, machine_pow_j, shared_params, model_state_i, model_state_j, sigma_i, sigma_j)

Estimate pairwise fidelity-like overlap and its gradient w.r.t shared parameters.

This reuses the same estimator structure as the current multi-state fidelity_expect_and_grad_joint kernel, but treats the two heads as tied copies of one parameter pytree. The gradient is computed by differentiating a surrogate f(p_i, p_j) and then evaluating on the diagonal p_i = p_j = shared_params. The exact gradient w.r.t the shared parameters is

\[\nabla_\theta f(\theta, \theta) = \partial_{p_i} f + \partial_{p_j} f.\]

Therefore we return grads['i'] + grads['j'] as the correct shared-parameter penalty gradient.

Notes

• For machine_pow == 2 in both samplers this estimates the normalized quantum fidelity.

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.

property states: list[MCState]

estimate(observables)

Return MCMC statistics for the expectation value of observables in the current state of the driver.

Parameters:

observables – A pytree of operators for which statistics should be computed.

Returns:

A pytree of the same structure as the input, containing MCMC statistics for the corresponding operators as leaves.