Ansätze

Restricted Boltzmann Machine (RBM)
- Real Restricted Boltzmann Machine (real-RBM)
- Complex Restricted Boltzmann Machine (complex-RBM)
- Cosine Restricted Boltzmann Machine (cos-RBM)
- Tanh Restricted Boltzmann Machine (tanh-RBM)
- Phase Restricted Boltzmann Machine (phase-RBM)
- Autoregressive Restricted Boltzmann Machine (AR-RBM)
- Ising-type Restricted Boltzmann Machine (Ising-RBM)
- Restricted Ising-type Restricted Boltzmann Machine (RIsing-RBM)
Recurrent Neural Network (RNN)
- Recurrent Neural Network (RNN)
- Gated Recurrent Unit (GRU)
- Graph MPS(Tensor)–RNN (MPS(Tensor)–RNN)
Transformer
Mix-Ansatz

RBM

real(complex)-RBM

\[\begin{split}\begin{split} \psi_{\theta}(n) & = \textcolor{teal}{\exp}{\sum_{j=1}^{N_{\rm v}}a_jn_j} \times \prod_i^{N_{\rm h}}\textcolor{violet}{2\cosh}(b_i + \sum_{j=1}^{N_{\rm v}}W_{ij}n_j) \\ \text{or} & = \prod_i^{N_h}\textcolor{violet}{2\cos}(b_i + \sum_{j=1}^{N_{\rm v}}W_{ij}n_j) \quad \textbf{cos-type}\\ \text{or} & = \textcolor{teal}{\tanh}{\sum_{j=1}^{N_{\rm v}}a_jn_j} \times \prod_i^{N_{\rm h}}\textcolor{violet}{2\cosh}(b_i + \sum_{j=1}^{N_{\rm v}}W_{ij}n_j) \quad \textbf{tanh-type} \end{split}\end{split}\]

For more information, see: ./vmc/ansatz/multi/RBMWavefunction.

Transformer

use nano-chatgpt

For more information, see: ./vmc/ansatz/transformer/decoder/DecoderWaveFunction.

MPS-RNN

For more information, see: ./vmc/ansatz/rnn/graph_mpsrnn/Graph_MPS_RNN.

Mix-Ansatz

Define: \(\psi(n) = f_n\phi(n), \ket{n} \sim |\phi(n)|^2\). \(\phi(n)\) is MPS-RNN, Transformer, AR–RBM with \(|\phi(n)|^2=1\) for sampling, \(f_n\) is RBM, MLP, Jastrow Factor, Transformer and so on.

\[\begin{split}\begin{aligned} B & = \left\langle |f_n|^2\right\rangle_{n \sim{|\phi(n)|^2} } \\ \widetilde{f}_n & = f_n /\sqrt{B} \\ E_{\rm loc}(n) &= \dfrac{\dfrac{\sum_m f_n^* H_{nm}f_m\phi(m)}{\phi(m)}}{\langle |f_n|^2\rangle} = \dfrac{\sum_m \widetilde{f}_n^* H_{nm}\widetilde{f}_m\phi(m)}{\phi(n)} \\ \partial_\theta \langle E\rangle &= 2\Re\big\langle (\partial_\theta (\ln(f_n\phi(n)))^*)(E_{\rm loc}(n) - \langle E\rangle|\widetilde{f}_n|^2) \big\rangle_{n\sim |\phi(n)|^2} \\ \end{aligned}\end{split}\]

Penalty NQS common-envelope sampling

Penaltywavefunctions(..., use_1sample=True) samples the positive amplitude \(\sqrt{\rho(n)}\), where \(\rho(n)=\sum_i|\Psi_i(n)|^2\), once per iteration. The default False retains one independent sampler per state. Both MCMC and Exact sampling are supported. Exact sampling requires ExactParams.alpha=2; MCMC uses its usual reweighting back to exponent two (gamma=None or gamma=2-alpha).

All states use the same configurations. If \(f_a\) is the envelope sample frequency (or the exact envelope probability), their weights are

\[b_a=\frac{f_a}{\rho(n_a)},\qquad D_i=\sum_a b_a|\Psi_i(n_a)|^2,\qquad P_{ia}=\frac{b_a|\Psi_i(n_a)|^2}{D_i}.\]

The overlap estimator and penalty are

\[B_{ij}=\sum_a b_a\Psi_i^*(n_a)\Psi_j(n_a),\qquad R_{ij}=\frac{B_{ij}}{D_i},\qquad F_{ij}=R_{ij}R_{ji}=\frac{|B_{ij}|^2}{D_iD_j}.\]

This common-sample Gram matrix keeps \(F_{ij}\) in \([0,1]\) when the state norms are resolved. Individual states may vanish on another state’s samples: local kernels evaluate \(H\Psi_i\) directly, and gradient and SR/minSR construction use amplitude derivatives at these nodes. No samples are discarded. The importance reference is fixed during differentiation; state-normalization derivatives are retained. use_twoside, use_diag and use_subspace retain their meanings and existing restrictions. Distributed weights are normalized across all ranks.

For Exact enumeration the envelope factor cancels. Both settings therefore use the same direct ONV sums and amplitude derivatives, with \(b_a=1\). Exact minSR evaluates the same damped solution by SVD of the weighted Jacobian, avoiding the loss of precision from forming its sample-space Gram matrix.

The log entry state i energy reports e_mean[i] + ecore in the original model-state order. The separate subspace root i eigenvalue entries report the sorted generalized eigenvalues e_diag[i] + ecore. The root index can therefore differ from the model-state index when their energies cross.

Multistate VAFQMC Ansatz

MultistateVAFQMCAnsatz contains \(K\) VAFQMC branches \(\{\psi_i\}_{i=0}^{K-1}\). The branches are sampled with one common auxiliary-field envelope, while the energy is evaluated from the cross-state overlap and Hamiltonian matrices

\[S_{ij} = \langle \psi_i | \psi_j \rangle, \qquad H_{ij} = \langle \psi_i | \hat H | \psi_j \rangle .\]

The multistate roots are obtained from the generalized eigenvalue equation

\[H c_k = E_k S c_k .\]

For numerical stability, the sampled Hamiltonian can be shifted by a reference energy \(E_{\rm ref}\):

\[H' = H - E_{\rm ref} S, \qquad E_k = E'_k + E_{\rm ref}.\]

The normalized overlap matrix is

\[\widetilde S_{ij} = \frac{S_{ij}}{\sqrt{S_{ii}S_{jj}}} .\]

The smallest eigenvalue of the Hermitian part of \(\widetilde S\) is required to be larger than Smin. This checks whether the sampled branches remain resolved; the matrices used for the generalized eigenproblem are not Hermitianized. Roots with non-negligible imaginary parts reject the update. Each logged step reports the L2 difference between each matrix used for the root solve and its Hermitian part; this is diagnostic only.

The optional sign/phase penalty is based on

\[q_i = \frac{S_{ii}}{\left\langle |S_{ii}(x)| \right\rangle_x}, \qquad L_{\rm sign} = \frac{1}{K}\sum_i \left[\max(q_{\rm target}-q_i,0)\right]^2 .\]

Its contribution to the objective is \(\beta L_{\rm sign}\). The pairwise normalized-overlap penalty is

\[L_{\rm overlap} = \operatorname{mean}_{i<j} \left[\max\left(|\widetilde S_{ij}|-s_{\rm target},0\right)\right]^2,\]

with contribution \(\alpha L_{\rm overlap}\). The total multistate objective is therefore

\[L = \frac{1}{K}\sum_k E_k + \beta L_{\rm sign} + \alpha L_{\rm overlap} .\]

The constructor arguments are:

states: The branch ansatzes. All branches must have the same field shape, orbital count, and electron count.
e_ref: The reference energy \(E_{\rm ref}\) used in the shifted Hamiltonian.
use_real: If true, solve the generalized eigenproblem using real matrices; otherwise retain the complex sampled matrices. Neither case Hermitianizes them.
Smin: Lower bound for the smallest eigenvalue of the Hermitian part of \(\widetilde S\).
k_max: Upper bound for the sampled absolute overlap kernel \(|S_{ij}(x)/r(x)|\).
beta and beta_t: Weight and target of the branch sign/phase penalty.
alpha and alpha_t: Weight and allowed normalized pair overlap for the branch-overlap penalty.

Penalty VAFQMC Ansatz

PenaltyVAFQMCAnsatz uses the same branch states, auxiliary-field envelope, and sampled \(H\) and \(S\) matrices as MultistateVAFQMCAnsatz. Its training objective is

\[\widetilde S = \tfrac12(S+S^\dagger), \qquad E_i = E_{\rm ref} + \frac{H'_{ii}}{S_{ii}}, \qquad L = \sum_i w_i E_i + \sum_{i>j}\alpha_{ij} \frac{|\widetilde S_{ij}|^2}{S_{ii}S_{jj}} + \beta L_{\rm sign}.\]

Here alpha is a finite symmetric matrix with a zero diagonal; energy_weights defaults to ones. beta and beta_t retain the multistate sign penalty defined above. For real-state runs, \(S\) is replaced by \(\operatorname{Re}S\) in the penalty. use_twoside=False (the default) assigns each pair’s overlap-penalty gradient only to the higher-numbered state; True differentiates it through both states. The logged loss is the same in either mode. Generalized eigenvalues of the sampled \(H\) and \(S\) are also logged each step as diagnostics, without entering the penalty objective. Hermitianization only makes the finite-sample overlap penalty nonnegative; the sampled matrices are not Hermitianized before solving. The logged \(H\) and \(S\) use unit-diagonal branch normalization, as in NQS penalty; this scaling does not change the generalized eigenvalues.

For example, with three existing branch ansatzes:

ansatz = PenaltyVAFQMCAnsatz(
    states,
    alpha=[[0, 1, 1], [1, 0, 1], [1, 1, 0]],
    energy_weights=[1, 1, 1],
    use_twoside=False,
    e_ref=reference_energy,
    beta=1.0,
)

Pass ansatz to VAFQMCOptimizer as for the multistate example. Set intermediate=[500, 1000, 2000] on the optimizer to save separate full checkpoints after those completed updates as <prefix>-intermediate-500.pth, <prefix>-intermediate-1000.pth, and <prefix>-intermediate-2000.pth. Each contains model for a later AFQMC trial and step for the completed update count; epoch remains zero-based for existing checkpoint readers. The interval option controls the separately overwritten <prefix>-checkpoint.pth file.

For Wheeler’s weighted two-sided objective, use:

ansatz = PenaltyVAFQMCAnsatz.wheeler(
    states,
    energy_weights=[3, 2, 1],
    alpha=0.75,
    e_ref=reference_energy,
    beta=1.0,
)

This entry point requires positive, strictly decreasing weights and a positive scalar alpha. It sets every off-diagonal pair coefficient to alpha and always uses two-sided gradients. Neither weights nor coefficients are rescaled. Each unordered pair is counted once. With exact overlaps, the unregularized Wheeler objective targets the individual eigenstates when

\[\alpha > \max_{i<j} (E_j-E_i)\frac{w_iw_j}{w_i-w_j}.\]

See Wheeler et al., Electronic Structure 6, 025001 (2024), doi:10.1088/2516-1075/ad38f8. This exact-objective condition is not a convergence guarantee for finite samples, restricted ansatzes, or the additional VAF sign penalty. The existing raw sampled overlap estimator is retained. Evaluate the original branch energies and their overlaps with reference eigenstates; diagonalizing the branch space is not part of this objective and must not be used to claim that each branch has converged.

Spin-flip

see: branch spin-flip

\[\begin{split}\begin{aligned} B & = \bigg\langle |f_n|^2 + \eta f^*_n f_{\bar n }\frac{\phi(\bar n)}{\phi(n)}\bigg\rangle_{n \sim{\phi_n^2} } \\ \widetilde{f}_n & = f_n /\sqrt{B} \\ E_{\rm loc}(n) &= \frac{\sum_m \widetilde{f}_n^* H_{nm} (\widetilde{f}_m\phi_m + \eta\widetilde{f}_{\bar m}\phi_{\bar m})} {\phi_n} \\ C & = \frac{|f_n|^2 + \eta f^*_n f_{\bar n }\frac{\phi(\bar n)}{\phi(n)}}{B} \\ \partial_\theta E &= 2\Re\left< (\partial_\theta (\ln(\phi_n f_n))^*)(E_{\rm loc}(n) - \left\langle E \right\rangle C) \right> \end{aligned}\end{split}\]