neuralqx.experimental.solver package

Experimental end-to-end solvers in neuraLQX.

class MultiSolver(lqx, output_path=None, auxiliary_path=None, *, clean_up=False, seed=None)

Bases: Solver

Solver subclass that trains multiple variational states jointly.

This solver is a drop in variant of Solver that runs MultiStateVMC with a MultiMCState backend. It supports configuring multiple networks, building one MCState per network during VMC initialization, evaluating expectations per state or for all states, exporting and importing multi state checkpoints, and logging final constraint results both aggregated and per state.

Returns:

None.

class STMultiSolver(lqx, output_path=None, auxiliary_path=None, *, clean_up=False, seed=None)

Bases: Solver

Solver for the single-trunk multi-head (ST-MH) variational ansatz.

Compared to MultiSolver (MT-MH / independent networks), this solver builds one shared multi-head Flax model and exposes one MCState per head using head-selector wrappers. The variational state is a STMultiMCState, and optimization is performed by SingleTrunkMultiHeadVMC, which aggregates energy and orthogonality gradients into a single shared parameter update.

Submodules