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.

Submodules