quantax.state.RestrictedDetState#

class quantax.state.RestrictedDetState#

Bases: MeanFieldFermionState

Restricted determinant mean-field state, a wrapper of RestrictedDet.

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.

rho_from_model = Partial( func=_JitWrapper( fn='RestrictedDetState.rho_from_model', filter_warning=False, donate_first=False, donate_rest=False ), args=(quantax.state.fermion_mf.RestrictedDetState,), keywords={} )#

property Nmodes: int#: Number of modes

property Nparticles: int | tuple[int, int] | None#: Number of particle convervation of the state

property Nsites: int#: Number of sites

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

property basis#: Quspin basis of the 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.

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

property energy: Array | ndarray | bool | number | bool | int | float | complex | None#: The energy in the previous optimization step.

expectation(operator: Operator, samples: Samples | jax.Array, return_var: bool = False) → complex | tuple[complex, float]#

The expectation value \(\left< \psi|O|\psi \right> / \left< \psi|\psi \right>\) of an operator, estimated from the given samples. This is a convenience wrapper of expectation.

Parameters:

operator – The operator \(O\).
samples – The samples for estimating the expectation value.
return_var – Whether the variance should also be returned, default to False.

Returns:

The mean value, and additionally the variance when return_var = True.

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\)

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

get_loss_fn(hamiltonian: Operator)#

Get the loss function for optimization.

Parameters:: hamiltonian – The Hamiltonian to compute the gradient of.

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.

get_step(hamiltonian: Operator) → jax.Array#

Get the gradient of the energy with respect to the mean-field parameters.

Parameters:: hamiltonian – The Hamiltonian to compute the gradient of.

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

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

classmethod is_paired() → bool#: Whether the state is a paired state (pfaffian) or not (determinant), default to False

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.

mf_expectation(operator: Operator) → jax.Array#

Compute the expectation value of an operator.

Parameters:: operator – The operator to compute the expectation value of. It should be an instance of Operator.

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

norm(ord: int | None = None) → Array#

Norm of state, returned as a real scalar.

Parameters:: ord – Order of the norm, default to 2-norm \(\sqrt{\sum_s |\psi(s)|^2}\)

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

property nsymm: int#: Number of symmetry group elements

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.

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

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.

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

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

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)\).

property symm: Symmetry#: Symmetry of the state

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.

todense(symm: Symmetry | None = None) → DenseState#

Obtain the quantax.state.DenseState corresponding to the current state

Parameters:: symm – The symmetry of the state, default to the current symmetry of the state

Warning

Users are responsible to ensure that the state satisfies the given symm.

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\).

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

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