Signal to noise Ratio regularization for tdvp

Created by: danielalcalde

I have at the moment hacky implementation of the SNR regularization from https://arxiv.org/pdf/1912.08828.pdf, which for the system I'm studying leads to a more stable TDVP evolution.

To compute the SNR you need to compute both the mean and variance of \rho_k(s)=(V^\dagger) _{k, k’} O_k’(s) O_loc(s). Where V are the eigenvectors of the quantum geometric tensor. The SNR can then be computed with snr_k=sqrt(N_samples) * abs(\rho_mean_k)/\rho_std_k.

The last step of the algorithm is to remove all \rho_k with an SNR smaller than a threshold (2 works quite well). \rho_mean_k is just a linear transformation of the F_k vector, but to compute the variance you need the O_k(s)*O_loc(s).

The main challenge to get this working cleanly in netket is to somehow pass O_k(s)*O_loc(s) to the solver. I suggest adding a new TDVP equation class that could then be easily be customized for this or other purposes.

The idea would be that instead of doing this:

E, F = variational_state.expect_and_grad(hamiltonian)
S = variational_state.quantum_geometric_tensor(qgt_T) + diag_shift
dw, _ = S.solve(solver, F)

we would be doing something like this:

tdvp_eq = variational_state.tdvp_equation(qgt_T)
dw, _ = tdvp_eq.solve(solver, diag_shift=diag_shift)

I'm open to suggestions.