netket_foundation.FoundationalQuantumState#
- class netket_foundation.FoundationalQuantumState[source]#
Bases:
VariationalStateVariational state that jointly samples physical configurations and fixed parameter replicas.
- Attributes
-
- chunk_size[source]#
Suggested maximum size of the chunks used in forward and backward evaluations of the Neural Network model.
If your inputs are smaller than the chunk size this setting is ignored.
This can be used to lower the memory required to run a computation with a very high number of samples or on a very large lattice. Notice that inputs and outputs must still fit in memory, but the intermediate computations will now require less memory.
This option comes at an increased computational cost. While this cost should be negligible for large-enough chunk sizes, don’t use it unless you are memory bound!
This option is an hint: only some operations support chunking. If you perform an operation that is not implemented with chunking support, it will fall back to no chunking. To check if this happened, set the environment variable NETKET_DEBUG=1.
- n_samples_per_rank[source]#
The number of samples generated on every jax device at every sampling step.
- samples[source]#
Returns the set of cached samples.
The samples returned are guaranteed valid for the current state of the variational state. If no cached parameters are available, then they are sampled first and then cached.
To obtain a new set of samples either use
reset()orsample().
- sampler_state: SamplerState[source]#
The current state of the sampler.
- Methods
- expect(O)[source]#
Estimates the quantum expectation value for a given operator \(O\) or generic observable. In the case of a pure state \(\psi\) and an operator, this is \(\langle O\rangle= \langle \Psi|O|\Psi\rangle/\langle\Psi|\Psi\rangle\) otherwise for a mixed state \(\rho\), this is \(\langle O\rangle= \textrm{Tr}[\rho \hat{O}]/\textrm{Tr}[\rho]\).
- Parameters:
O (
AbstractOperator) – the operator or observable for which to compute the expectation value.- Return type:
- Returns:
An estimation of the quantum expectation value \(\langle O\rangle\).
- expect_and_forces(O, *, mutable=None)[source]#
Estimates the quantum expectation value and the corresponding force vector for a given operator O.
The force vector \(F_j\) is defined as the covariance of log-derivative of the trial wave function and the local estimators of the operator. For complex holomorphic states, this is equivalent to the expectation gradient \(\frac{\partial\langle O\rangle}{\partial(\theta_j)^\star} = F_j\). For real-parameter states, the gradient is given by \(\frac{\partial\partial_j\langle O\rangle}{\partial\partial_j\theta_j} = 2 \textrm{Re}[F_j]\).
- Parameters:
O (
AbstractOperator) – The operator O for which expectation value and force are computed.mutable (
Union[bool,str,Collection[str],DenyList,None]) – Can be bool, str, or list. Specifies which collections in the model_state should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. This is used to mutate the state of the model while you train it (for example to implement BatchNorm. Consult Flax’s Module.apply documentation for a more in-depth explanation).
- Return type:
- Returns:
An estimate of the quantum expectation value <O>. An estimate of the force vector \(F_j = \textrm{Cov}[\partial_j\log\psi, O_{\textrm{loc}}]\).
- expect_and_grad(O, *, mutable=None, **kwargs)[source]#
Estimates the quantum expectation value and its gradient for a given operator \(O\).
- Parameters:
O (
AbstractOperator) – The operator \(O\) for which expectation value and gradient are computed.mutable (
Union[bool,str,Collection[str],DenyList,None]) –Can be bool, str, or list. Specifies which collections in the model_state should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. This is used to mutate the state of the model while you train it (for example to implement BatchNorm. Consult Flax’s Module.apply documentation for a more in-depth explanation).
use_covariance – whether to use the covariance formula, usually reserved for hermitian operators, \(\textrm{Cov}[\partial\log\psi, O_{\textrm{loc}}\rangle]\)
- Return type:
- Returns:
An estimate of the quantum expectation value <O>. An estimate of the gradient of the quantum expectation value <O>.
- get_state(parameters, seed=None)[source]#
Given a set of parameters, returns a standard MCState instance that corresponds to the foundational state with those parameters.
- grad(Ô, *, use_covariance=None, mutable=None)[source]#
Estimates the gradient of the quantum expectation value of a given operator O.
- Parameters:
op (
netket.operator.AbstractOperator) – the operator O.is_hermitian – optional override for whether to use or not the hermitian logic. By default it’s automatically detected.
use_covariance (bool | None)
mutable (bool | str | Collection[str] | DenyList | None)
- Returns:
An estimation of the average gradient of the quantum expectation value <O>.
- Return type:
- init(seed=None, dtype=None)[source]#
Initialises the variational parameters of the variational state.
- init_parameters(init_fun=None, *, seed=None)[source]#
Re-initializes all the parameters with the provided initialization function, defaulting to the normal distribution of standard deviation 0.01.
Warning
The init function will not change the dtype of the parameters, which is determined by the model. DO NOT SPECIFY IT INSIDE THE INIT FUNCTION
- Parameters:
init_fun (
Callable[[Any,Sequence[int],Union[None,str,type[Any],dtype,_SupportsDType]],Array] |None) – a jax initializer such asjax.nn.initializers.normal(). Must be a Callable taking 3 inputs, the jax PRNG key, the shape and the dtype, and outputting an array with the valid dtype and shape. If left unspecified, defaults tojax.nn.initializers.normal(stddev=0.01)seed (
Any|None) – Optional seed to be used. The seed is synced across all JAX processes. If unspecified, uses a random seed.
- is_state(target_parameters, *, reference=None, chunk_size=None)[source]#
Create an
ISStatetargetingtarget_parameters.By default the current physical samples are used as the IS reference, with log-probabilities computed from the joint samples (which carry each replica’s foundational parameters), so IS weights correctly account for the per-replica distribution.
Alternatively, pass
referenceto use a precomputed reference distribution — either a path to a.npzfile written bysave(), aSamplesWithProb, or a raw(samples, log_probs)tuple.- Parameters:
target_parameters (
Array) – 1-D array of foundational parameters for the target state.reference – Optional IS reference: a
.npzpath, aSamplesWithProb, or a(samples, log_probs)tuple. Defaults to the current physical samples.chunk_size (
int|None) – Forwarded to ISState; defaults to self.chunk_size.
- Returns:
ISState ready for .expect().
- log_value(σ)[source]#
Evaluate the variational state for a batch of states and returns the logarithm of the amplitude of the quantum state.
For pure states, this is \(\log(\langle\sigma|\psi\rangle)\), whereas for mixed states this is \(\log(\langle\sigma_r|\rho|\sigma_c\rangle)\), where \(\psi\) and \(\rho\) are respectively a pure state (wavefunction) and a mixed state (density matrix). For the density matrix, the left and right-acting states (row and column) are obtained as
σr=σ[::,0:N]andσc=σ[::,N:].Given a batch of inputs
(Nb, N), returns a batch of outputs(Nb,).
- quantum_geometric_tensor(qgt_T)[source]#
Computes an estimate of the quantum geometric tensor G_ij.
This function returns a linear operator that can be used to apply G_ij to a given vector or can be converted to a full matrix.
- Parameters:
qgt_T (
Callable[[VariationalState],LinearOperator] |None) – the optional type of the quantum geometric tensor. By default it is automatically selected.- Returns:
- A linear operator representing the quantum
geometric tensor.
- Return type:
LinearOperator
- replace_sampler_seed(seed=None)[source]#
Change the rng state of the sampler contained in this Monte Carlo state.
The use-case for this is when you create a copy of a state (for example by loading it from disk), but you want the copy to generate samples that are not correlated to the ones of the original state.
Beware that for this to work correctly, you probably need to resample a bunch of times in order to decorrelate the chains, because this method only changes the rng seed, but not the current configurations in the chain.
- reset()[source]#
Resets the sampled states. This method is called automatically every time that the parameters/state is updated.
- sample(*, chain_length=None, n_samples=None, n_discard_per_chain=None)[source]#
Sample a certain number of configurations.
If one among chain_length or n_samples is defined, that number of samples are generated. Otherwise the value set internally is used.