quantax.model.GeneralJastrow

class quantax.model.GeneralJastrow

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.

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

Evaluate the wavefunction on a given configuration.

property use_ref: bool

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

property fermion_mf: RefModel

The reference state \(\psi_\mathrm{mf}\).

get_sublattice_spins(x: Array) → Array

Return the configurations translated over the sublattice, one per translation kept in the symmetrization.

sub_symmetrize(x_net: NDArray | Array | LogArray | ScaleArray, x_mf: NDArray | Array | LogArray | ScaleArray, s: Array) → NDArray | Array | LogArray | ScaleArray

Combine the network output x_net with the reference amplitudes x_mf evaluated on the translated configurations into the symmetrized wavefunction.

init_internal(s: Array) → tuple[NDArray | Array | LogArray | ScaleArray, MF_Internal]

Return wavefunction and internal values for given input configurations. See RefModel for details.

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: MF_Internal, return_update: Literal[False] = False) → NDArray | Array | LogArray | ScaleArray

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: MF_Internal, return_update: Literal[True]) → tuple[NDArray | Array | LogArray | ScaleArray, MF_Internal]

Accelerated forward pass through local updates and internal quantities. The reference state is updated through its own ref_forward while the network factor is recomputed from scratch. See RefModel for details.