neuralqx.solver package

This module implements end-to-end solvers for neuraLQX.

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

Bases: ABC

Abstract base class for neuraLQX solvers.

A solver is responsible for orchestrating the full optimisation workflow around a given LQX model: configuration, initialisation, execution, logging, checkpointing, and result visualisation.

This base class provides: - consistent wiring of model components (graph, gauge group, Hilbert interface), - reproducible run identity via a unique hash and seed, - output directory management with MPI-safe directory creation and optional cleanup, - a unified logging surface via a runtime log object, - a set of “initialisation flags” that gate access to derived objects (sampler, optimiser, network,

driver, variational state).

Concrete implementations must provide the operational methods such as: - constructing samplers/optimisers/networks, - building the variational state and VMC driver, - running and resuming optimisation loops, - exporting/importing state, - plotting results and observables.

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

Bases: AbstractSolver

Concrete neuraLQX solver implementing a VMC-based optimisation workflow.

This class is the main user-facing solver implementation that executes variational optimisation for an LQX model. It inherits core wiring, output/log bookkeeping, and safety flags from AbstractSolver and implements the operational API:

  • set_sampler() builds and attaches a NetKet-compatible sampler (including specialised gauge samplers when supported by the Hilbert space).

  • set_optimizer() builds an Optax optimiser and configures stochastic reconfiguration (SR), including the SR linear solver backend.

  • set_network() installs a neural network Ansatz and optionally wraps it with symmetry/group projection when running diffeomorphism-invariant simulations.

  • initialize_vmc() constructs the MCState and the VMC driver, including SR preconditioning.

  • run() executes an optimisation loop with robust interruption handling, finalisation, state export, and plotting.

  • continue_simulation() resumes from an in-memory or on-disk checkpoint while preserving runtime logs and ensuring driver/state consistency.

Key design goals

Robustness and reproducibility:

  • Every solver run has a unique hash and a deterministic seed (unless user overrides).

  • State export/import is MPI-aware and uses broadcast to ensure all ranks reconstruct the same state.

  • Finalisation is defensive: I/O and plotting failures do not crash the solver after a successful run; they are treated as best-effort.

Clarity of run semantics:

  • The run pipeline is explicitly split into preparation, execution, and finalisation phases.

  • Continuations do not implicitly reset logs or state unless explicitly requested.

MPI-safe interruption:

  • Single-rank runs attempt graceful abort and still export a final checkpoint.

  • Multi-rank runs hard-abort on Ctrl+C to avoid deadlocks from rank-skewed collectives.

type lqx:

AbstractLqxInterface

param lqx:

The model interface implementing AbstractLqxInterface.

type output_path:

str | None

param output_path:

Optional base directory for solver outputs.

type auxiliary_path:

str | None

param auxiliary_path:

Optional grouping component appended under output_path.

type clean_up:

bool | None

param clean_up:

If True, remove empty stale directories before starting.

type seed:

int | None

param seed:

Optional explicit seed controlling solver reproducibility.

return:

None.

raises ValueError:

If output_path is provided but invalid.

raises OSError:

If output directories cannot be created.

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.

serialize_MultiMCState(vstate)

Serialize a MultiMCState into a plain Python dictionary.

The resulting dictionary is suitable for writing to disk with the neuraLQX serialization utilities. Each contained MCState is serialized using serialize_MCState.

Parameters:

vstate (MultiMCState) – Multi state variational state to serialize.

Return type:

dict

Returns:

Dictionary representation of the multi state, including per state payloads.

deserialize_MultiMCState(template, state_dict, *, force_load_mpi=False)

Deserialize a MultiMCState from a dictionary into a new MultiMCState instance.

The function uses an existing template multi state to provide structure and implementation details required by deserialize_MCState. For backward compatibility, a dictionary that represents a single MCState is accepted and wrapped into a single state MultiMCState.

Parameters:

  • template (MultiMCState) – Existing multi state used as a template for reconstructing each state.

  • state_dict (dict) – Dictionary payload previously produced by serialization utilities.

  • force_load_mpi (bool) – If True, attempt to load MPI related fields even if they differ from the current run configuration.

Return type:

MultiMCState

Returns:

Newly reconstructed MultiMCState.

Raises:

ValueError – If the number of saved states does not match template.n_states.

Submodules