SPRING (Kaczmarz) optimizer

[wdphy16]

The discussion started in #1696 . It's particularly interesting to me because it may fulfill my wish for years to properly have a running estimation of S^-1, so I hurried this while away for a meeting, and let's see if it really works well in our use cases.

We need to test if it fully supports complex numbers and MPI, and also discuss the API.

I added the arguments proj_reg and momentum to the current API of VMC_SRt. When momentum = 0, it degenerates to the previous SRt, and the matrix P will not affect the training as long as the training was already numerically stable.

Here I call the matrix P 'projector regularization', or maybe we can call it 'ones regularization'. Multiplying it by any factor will not change the updates, as long as it's numerically stable, even if we split the complex parameters. I've tested it numerically.

A N * N matrix of ones has only one nonzero eigenvalue, which is N. To make it not too large, we divide it by N by default. In principle the user can adjust the factor proj_reg but I guess it's not needed in most cases.

The momentum subtracted from dv and added to updates are implemented in the preconditioner. For the gradient norm clipping, however, we can use optax to do it after the preconditioner.

By the way, after we clean up the driver API (e.g. in #1674 ), maybe it's a good time to factor out SRt to be actually a preconditioner.

cc @attila-i-szabo @riccardo-rende @llviteritti

[codecov]

Codecov Report

Attention: 4 lines in your changes are missing coverage. Please review.

Comparison is base (a4a9af2) 82.71% compared to head (13576b4) 51.23%.

Files	Patch %	Lines
netket/experimental/driver/vmc_srt.py	80.00%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1708       +/-   ##
===========================================
- Coverage   82.71%   51.23%   -31.48%     
===========================================
  Files         298      298               
  Lines       18236    18250       +14     
  Branches     3501     3504        +3     
===========================================
- Hits        15083     9351     -5732     
- Misses       2478     8214     +5736     
- Partials      675      685       +10     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[PhilipVinc]

I will find the time to properly review the code later in the week, however there's a comment I would like to make first.

In general, to keep the maintainance effort reasonable, I try to keep within netket only 'stuff' that is used by a reasonable number of persons and that is of proven utility.
I am somewhat against adding the latest 'hot' thing because who knows if it will be useful 1 year from now.

So, in this case, as this optimiser was published just a few weeks ago, I feel obliged to ask the question: is this worth implementing within netket? is it really so much better?
Can't we keep it in a separate repository a la netket fidelity for the time being, and if we see that it is really great, we merge it here?