quantax.state.Variational#

class quantax.state.Variational#

Bases: State

Variational state. This is a wrapper of a jittable variational ansatz. The variational model should be given as an equinox.Module. For details of Equinox, see this documentation.

Warning

There are many intermediate values stored in the class, so most functions like __call__ in this class are non-jittable.

Warning

Many quantities are only computed once in the initialization. Please don’t change the private attributes unless an update function is provided. One can define a new state if some changes are necessary.

Parameters:

model – Variational model. Should be an equinox.Module.
param_file – File for loading parameters which is saved by save or equinox.tree_serialise_leaves, default to not loading parameters.
symm – Symmetry of the network, default to Identity. Denoting the network output as \(f(s)\) and symmetry elements as \(T_i\) with characters \(\omega_i\), the wave function is given by \(\psi(s) = \sum_i \omega_i \, f(T_i s) / n_{symm}\)
max_parallel –
The maximum chunk size allowed per device. Specifying a limited value is important for avoiding memory overflow. For many hamiltonians, this also helps to improve the efficiency by keeping a constant amount of forward pass and avoiding re-jitting. The allowed input formats are:
- None:
  No chunk size (default).
- int:
  The same chunk size for all forward and backward passes.
- Tuple[int, int]:
  (forward chunk, backward chunk)
- Tuple[int, int, int]:
  (forward chunk, backward chunk, ref chunk) Ref chunk is the chunk size used in init_internal and ref_forward. If ref chunk is not specified by this format, it defaults to the forward chunk size.
use_ref – Whether ref_forward and ref_forward_with_updates will be used when the model is a RefModel. When the model is not a RefModel, this argument has no effect. Default to True.

__call__(s: NDArray | Array) → NDArray | Array | LogArray | ScaleArray#

Evaluate the wavefunction \(\psi(s) = \left<s|\psi\right>\).

Parameters:: s – Spin/fermion configurations s with entries \(\pm 1\).

Warning

This function is not jittable.

Note

The returned value is \(\psi(s)\) instead of \(\log\psi(s)\).

property symm: Symmetry#: Symmetry of the state

property use_ref: bool#: Whether to use reference implementation for updates

property model: Callable[[Array], NDArray | Array | LogArray | ScaleArray]#: The variational model used in the variational state.

property holomorphic: bool#: Whether the variational state is holomorphic.

property forward_chunk: int | None#: The maximum chunk size of forward pass allowed per device.

property backward_chunk: int | None#: The maximum chunk size of backward pass allowed per device.

property ref_chunk: int | None#: The maximum chunk size of ref_forward allowed per device.

property nparams: int#: Number of total parameters in the variational state.

property dtype: str | type[Any] | dtype | SupportsDType#: The parameter data type of the variational state.

property vs_type: VS_TYPE#: The type of variational state.

fast_forward(s: Array) → NDArray | Array | LogArray | ScaleArray#

Evaluate the wavefunction \(\psi(s) = \left<s|\psi\right>\). This function assumes s to be in good shape and sharding for speedup.

Parameters:: s – Spin/fermion configurations s with entries \(\pm 1\)

init_internal(s: Array) → tuple[NDArray | Array | LogArray | ScaleArray, PyTree]#: Return the wavefunction and initial internal values for the given input s.

property required_update_modes: tuple[str, ...]#: The required update modes for accelerated ref_forward pass.

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: Literal[False] = False) → NDArray | Array | LogArray | ScaleArray#

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: Literal[True]) → tuple[NDArray | Array | LogArray | ScaleArray, PyTree]

Compute the forward pass given reference internal state of the model.

Parameters:

s – Input states s with entries \(\pm 1\).
s_old – The old states before the updates, with entries \(\pm 1\).
update_mode – The update modes required by the model.
internal – The internal state of the model, which is initialized by init_internal.

Returns:

A tuple of the output wave function \(\psi(s)\) and the updated internal state of the model.

segment_ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], idx_segment: Array, internal: PyTree) → NDArray | Array | LogArray | ScaleArray#

Compute the forward pass with segments given reference internal state of the model. This method is usually used in the computation of local energy.

Parameters:

s – Input states s with entries \(\pm 1\).
s_old – The old states before the updates, with entries \(\pm 1\).
update_mode – The update modes required by the model.
idx_segment – The indices of the segment to be updated, which is used to slice s_old and internal.
internal – The internal state of the model, which is initialized by init_internal.

Returns:

The output wave function \(\psi(s)\).

jacobian(s: Array) → Array#

Compute the jacobian matrix \(\frac{1}{\psi} \frac{\partial \psi}{\partial \theta}\). See VS_TYPE for the definition of jacobian for different kinds of networks.

Parameters:: s – The input spin/fermion configurations with entries \(\pm 1\).
Returns:: A 2D jacobian matrix with the first dimension for different inputs and the second dimension for different parameters. The order of parameters are the same as get_params_flatten.

partition(model: PyTree = None) → tuple[PyTree, PyTree]#

Split the variational model into two pytrees, one containing all parameters and the other containing all other elements, similar to partition in Equinox.

Parameters:: model – The model to be splitted, default to be the variational model in the variational state.

combine(params: PyTree, others: PyTree) → PyTree#

Combine two pytrees, one containing all parameters and the other containing all other elements, into one variational model. This is similar to combine in Equinox.

Parameters:

params – The pytree containing only parameters.
others – The pytree containing other elements.

get_params_flatten() → Array#: Obtain a flattened 1D array of all parameters.

get_params_unflatten(params: Array) → PyTree#: Obtain the parameters pytree from a flattened 1D array of all parameters.

update(step: Array) → None#

Update the variational parameters of the state as \(\theta' = \theta - \delta\theta\).

Parameters:: step – The update step \(\delta\theta\).

Note

The update direction is \(-\delta\theta\) instead of \(\delta\theta\).

save(file: str | Path | BinaryIO) → None#: Save the variational model in the given file. This file can be used be loaded when initializing Variational.

to_netket_model() → Module#

Convert the state to an equinox.Module compatible with NetKet. NetKet natively accepts equinox modules as the variational ansatz of an nk.vqs.MCState, so the returned module can be passed directly to NetKet to measure observables.

The module takes spin configurations with entries \(\pm 1\) (NetKet’s convention for netket.hilbert.Spin) and returns \(\log\psi\). Whatever the underlying quantax model outputs (jax.Array, LogArray, or ScaleArray), it is converted to \(\log\psi\). The output is always complex so that sign-structured or complex wavefunctions are represented correctly (see the wavefunction overflow section of the sharp_bits tutorial).

Warning

Training the generated state in NetKet is probably unstable, but the state can be reliably used to measure observables.