quantax.state.DenseState#

class quantax.state.DenseState#

Bases: State

Dense state in which the full wave function is stored as a numpy array

__init__(psi: NDArray | Array | LogArray | ScaleArray, symm: Symmetry | None = None)#

Parameters:

psi – Full wave function given according to the basis.states order in QuSpin
symm – The symmetry of the wave function, default to quantax.symmetry.Identity

__getitem__(basis_ints: NDArray | Array) → NDArray#

Evaluate the wave function \(\psi(s) = \left<s|\psi\right>\) by state[s]. This is done by slicing the full wave function.

Parameters:: basis_ints – A batch of basis integers

__call__(fock_states: NDArray | Array) → NDArray#

Evaluate the wave function \(\psi(s) = \left<s|\psi\right>\) by state(s). This is done by converting fock states basis integers and slicing the full wave function.

Parameters:: fock_states – A batch of fock states with entries \(\pm 1\)

property psi: NDArray | Array | LogArray | ScaleArray#: Full wave function

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

Convert the state to a new DenseState with the given symmetry

Parameters:: symm – The new symmetry. It’s default to self.symm, so self without copy is returned by default.

normalize() → DenseState#: Return a DenseState with normalized wave function.

normalize_() → None#: Normalize the wave function.

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

property dtype: str | type[Any] | dtype | SupportsDType#: The data type of the wave function values.

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

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 nsymm: int#: Number of symmetry group elements

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: bool = False) → NDArray | Array | LogArray | ScaleArray | tuple[NDArray | Array | LogArray | ScaleArray, PyTree]#: Compute the forward pass given reference internal state of the model.

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

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.

property symm: Symmetry#: Symmetry of the state

property use_ref: bool#: Whether to use reference implementation for local updates. Default to False.