quantax.model.GeneralJastrow

class quantax.model.GeneralJastrow(net: Callable[[jax.Array], PsiArray], fermion_mf: RefModel, trans_symm: Translation | None = None)

A reference state multiplied by a generalized (neural-network) Jastrow factor, \(\psi(s) = J(s)\, \psi_\mathrm{mf}(s)\), where \(J(s)\) is the output of an arbitrary network net and \(\psi_\mathrm{mf}\) is a reference state, usually a fermionic mean-field wavefunction.

When trans_symm is given, the factor is translation-symmetrized as \(\psi(s) = \sum_{T} \chi_T\, J_T(s)\, \psi_\mathrm{mf}(T s)\), where the sum runs over the translations \(T\), \(\chi_T\) are the corresponding characters, and \(J_T(s)\) is the network output associated with translation \(T\). If the reference state defines a sublattice, the sum is restricted to it accordingly.

__init__(net: Callable[[Array], NDArray | Array | LogArray | ScaleArray], fermion_mf: RefModel, trans_symm: Translation | None = None)

Parameters:

• net – The network producing the Jastrow factor \(J(s)\). To obtain a translation-resolved factor \(J_T(s)\), build the network with trans_symm=Identity() so that its output is not summed over translations.

• fermion_mf – The reference state \(\psi_\mathrm{mf}(s)\), typically a fermionic mean-field wavefunction. Its sublattice (if any) is reused to restrict the symmetrization.

• trans_symm – The translation symmetry used to symmetrize the product. None simply multiplies the (mean) network output with the reference state.

Methods

__init__(net, fermion_mf[, trans_symm])
get_sublattice_spins(x)	Return the configurations translated over the sublattice, one per translation kept in the symmetrization.
init_internal(s)	Return wavefunction and internal values for given input configurations.
ref_forward(...)	Accelerated forward pass through local updates and internal quantities.
sub_symmetrize(x_net, x_mf, s)	Combine the network output x_net with the reference amplitudes x_mf evaluated on the translated configurations into the symmetrized wavefunction.

Attributes

fermion_mf	The reference state \(\psi_\mathrm{mf}\).
required_update_modes	The required update modes for accelerated ref_forward pass.
use_ref	Whether to use reference implementation for local updates.
net
fermion_layer
holomorphic
trans_symm
sublattice