quantax.optimizer.SPRING#

class quantax.optimizer.SPRING#

Bases: StochasticQNGD

SPRING optimizer. This is a variant of SR with momentum.

The momentum makes large steps possible, so it is recommended to constrain the norm of the applied update in the training loop (see the note in MARCH).

__init__(state: Variational, hamiltonian: Operator, imag_time: bool = True, solver: Callable[[...], Array] | None = None, file: str | Path | BinaryIO | None = None, mu: float = 0.9, clip: float | None = 5.0)#

Initialize the SPRING optimizer.

Parameters:

  • state – Variational state to be optimized.

  • hamiltonian – The Hamiltonian for the evolution.

  • imag_time – Whether to use imaginary-time evolution.

  • solver – The numerical solver for the matrix inverse, default to auto_shift_eig.

  • file – The file with stored buffers of the optimizer.

  • mu – The momentum factor.

  • clip – Local-energy clipping passed to EnergyGrad, default to 5.0. Set to None to disable.

property VarE: Array | ndarray | bool | number | bool | int | float | complex | None#

Energy variance \(\left< (H - E)^2 \right>\) of the current step, None when the gradient source doesn’t define it.

property energy: Array | ndarray | bool | number | bool | int | float | complex | None#

Energy of the current step, None when the gradient source doesn’t define it.

get_Ebar(samples: Samples | Array) → Array#

Compute \(\bar \epsilon\) of the gradient source for given samples.

get_Obar(samples: Samples | Array) → Array#

Calculate \(\bar O = \frac{1}{\sqrt{N_s}}(\frac{1}{\psi} \frac{\partial \psi}{\partial \theta} - \left< \frac{1}{\psi} \frac{\partial \psi}{\partial \theta} \right>)\) for given samples.

get_step(samples: Samples | Array) → Array#

Obtain the optimization step by solving the equation \(\bar O \dot \theta = \bar \epsilon\) for given samples.

property hamiltonian: Operator | None#

The Hamiltonian for the evolution, None when the gradient source doesn’t define it.

property holomorphic: bool#

Whether the state is holomorphic.

property imag_time: bool#

Whether to use imaginary-time evolution.

save(file: str | Path | BinaryIO) → None#

Save the optimizer buffers to a file.

solve(Obar: Array, Ebar: Array, buffers: dict) → tuple[Array, dict]#

Generate the optimization step for given \(\bar O\) and \(\bar \epsilon\) by applying the updater strategy around solve_equation.

solve_equation(Obar: Array, Ebar: Array, **solver_kwargs) → Array#

Solve the linear equation \(\bar O \dot \theta = \bar \epsilon\) for given \(\bar O\) and \(\bar \epsilon\). Real and imaginary parts are stacked for non-holomorphic states, and extra keyword arguments are forwarded to the numerical solver.

A diag_preconditioner keyword is handled here: if the solver declares it as an argument (e.g. lsmr), it is passed through; otherwise the right preconditioning is emulated by solving with \(\bar O / d\) and rescaling the output by \(1 / d\).

property state: Variational#

Variational state to be optimized.

property vs_type: VS_TYPE#

The vs_type of the state, see VS_TYPE.