quantax.state.UnrestrictedDetState#

class quantax.state.UnrestrictedDetState#

Bases: MeanFieldFermionState

Unrestricted determinant mean-field state, a wrapper of UnrestrictedDet.

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

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

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

property Nmodes: int#: Number of modes

property Nparticles: int | tuple[int, int] | None#: Number of particle conservation 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) → Callable#

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 split, default to be the variational model in the variational state.

precompile(mesh: Mesh | None = None, forward_batches: Sequence[int] | int = (), backward_batches: Sequence[int] | int = ()) → None#

Ahead-of-time compile the expensive jitted functions of this state and write the executables into the persistent compilation cache.

Two usages:

precompile(...) without a mesh compiles for the devices of the current run, so the compilation cost is paid here instead of at the first forward/backward call later in the program (which then loads the executables from the persistent cache).
With a compile-only mesh from make_precompile_mesh, this compiles for a different (typically larger) device topology. A real run on that topology with the same environment loads the executables from the shared cache instead of compiling, so a cheap single-node job can absorb the compilation cost of a multi-node job.

Parameters:

mesh – The ("process", "device") mesh describing the target topology, default to make_mesh spanning the devices of the current run.
forward_batches – Global batch sizes to compile the jitted chunked forward (fast_forward) for; a single int is treated as one batch size. Jitted functions are compiled once per batch shape, and the batches of the target run are determined by its samplers: the chain numbers of the (sub-)samplers, and the chunk-multiple update batches of chunked sweeps. They can be read off a cheap pilot run of the target script with jax.config.update("jax_log_compiles", True) from the jit(chunked_f) log lines; the global batches are topology-independent, so a 1-GPU pilot with the target’s global number of samples reports the correct values.
backward_batches – Global batch sizes to compile the jacobian (jacobian) for, typically the global number of samples; a single int is treated as one batch size.

The forward used by __call__() and observable estimation is always compiled (its batch is fixed to forward_chunk * mesh.size by the chunking) when max_parallel is set. max_parallel must equal the value used by the target run.

Warning

The persistent-cache key includes the jaxlib version, XLA flags, GPU model, and CUDA version. Run the precompiling job on one node of the same cluster with the same environment and XLA_FLAGS as the target run, with jax_compilation_cache_dir pointing to a filesystem shared with it.

Note

Compiling for a compile-only mesh first executes a small real forward/backward warm-up on the current devices, with the target run’s per-device batches, to record cuDNN/GEMM autotuning results in the XLA per-fusion autotune cache; the ahead-of-time compilation reuses them instead of profiling (which is impossible without real devices). This requires the default jax_persistent_cache_enable_xla_caches setting, which keeps the autotune cache enabled alongside the persistent cache.

Note

Small helper functions (samplers, chunking utilities, solvers) still compile natively in the target run within seconds, and Python tracing is still paid by every process of the target run.

Note

The jitted functions of this state stay bound to the given mesh afterwards. With a compile-only mesh, executing them raises an error, so the state cannot be used for computation after precompiling for a different topology.

Note

Reference forwards (use_ref) are not covered yet.

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.
return_update – Whether to return the updated internal state along with the wave function.

Returns:

The output wave function \(\psi(s)\), or a tuple of the wave function and the updated internal state of the model if return_update is True.

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