quantax.state.State

class quantax.state.State

Abstract class for quantum states

__init__(symm: Symmetry | None = None)

Parameters:

symm – The symmetry of the state, default to quantax.symmetry.Identity

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

Evaluate the wave function \(\psi(s) = \left<s|\psi\right>\) by state(s)

Parameters:

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

__getitem__(basis_ints: NDArray | Array) → NDArray | Array | LogArray | ScaleArray

Evaluate the wave function \(\psi(s) = \left<s|\psi\right>\) by state[s]

Parameters:

basis_ints – A batch of basis integers

__matmul__(other: State) → complex

Compute the contraction \(\left< \psi|\phi \right>\) by self @ other. This is implemented by converting self and other to DenseState.

property Nsites: int

Number of sites

property Nmodes: int

Number of modes

property dtype: str | type[Any] | dtype | SupportsDType

The data type of the wave function values.

property symm: Symmetry

Symmetry of the state

property nsymm: int

Number of symmetry group elements

property basis

Quspin basis of the state

property Nparticles: int | tuple[int, int] | None

Number of particle convervation of the state

property use_ref: bool

Whether to use reference implementation for local updates. Default to False.

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.

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.

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.

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.

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

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

expectation(operator: Operator, samples: Samples | jax.Array, return_var: Literal[True]) → 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.