quantax.optimizer.AdamSR#

class quantax.optimizer.AdamSR#

Bases: StochasticQNGD

AdamSR optimizer. This is a variant of SR with first and second order momentum (like Adam). The time cost is roughly twice of SR.

A fixed learning rate leaves the parameters random-walking in a noise ball; for converged energies, decay the learning rate in the training loop, e.g. lr / (1 + max(t - t0, 0) / t0). Constraining the norm of the applied update (see MARCH) protects the ignition phase.

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

Initialize the AdamSR 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 first order momentum factor.

  • beta – The second order 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.