Density Matrix Improvements and Observables

[emerali]

• Vectorized Density Matrix Gradients

  • Cut run-time of density matrix tests on cpu (pytest -k "density and not gpu") from about 7m46s (running current state of master) down to about 37s 9s on my laptop.
  • There's still a bit more vectorization that can be done (ex. L377 of density_matrix.py), but this should be fine for now. EDIT: now that I think about it, should probably go through with this, as the density matrix tutorial takes a painfully long time.

• Created new base class NeuralStateBase, and moved a lot of code into there. Also removed some code that was duplicated across the wavefunctions and density matrix classes.

• Refactor observables to work with the new DensityMatrix object.

Todo: more testing, basically

• Use an Observable object to sample from a DensityMatrix. Make sure this gives the right values, since I'm using a different importance sampling scheme than what Giacomo did in his paper. I'm weighting by ratios of matrix elements of \rho, while in the paper they'd weight by the purified state.
• Make sure all the tutorials still run. Also double check the SciPost examples.
• Double check that SWAP operator is still working properly.
• Tweak the Observables tutorials a bit to show users how to write Observables classes that work for both pure and mixed states.
• Absorb density matrix versions of training_statistics into the main ones.
• Increase test coverage
• finish vectorizing density matrix gradients

[codecov-io]

Codecov Report

Merging #100 into master will increase coverage by 7.75%.
The diff coverage is 92.3%.

@@            Coverage Diff            @@
##           master    #100      +/-   ##
=========================================
+ Coverage   72.55%   80.3%   +7.75%     
=========================================
  Files          28      29       +1     
  Lines        1778    1559     -219     
  Branches      256     216      -40     
=========================================
- Hits         1290    1252      -38     
+ Misses        451     257     -194     
- Partials       37      50      +13

Impacted Files	Coverage Δ
qucumber/observables/observable.py	75.49% <ø> (+45.68%)	⬆️
qucumber/callbacks/metric_evaluator.py	41.17% <ø> (-17.65%)	⬇️
qucumber/observables/system.py	29.16% <ø> (ø)	⬆️
qucumber/callbacks/observable_evaluator.py	36.61% <ø> (ø)	⬆️
qucumber/observables/interactions.py	81.25% <ø> (+50%)	⬆️
qucumber/utils/data.py	79.48% <100%> (+26.65%)	⬆️
qucumber/nn_states/complex_wavefunction.py	86.95% <100%> (-5.67%)	⬇️
qucumber/observables/entanglement.py	100% <100%> (+77.77%)	⬆️
qucumber/observables/pauli.py	100% <100%> (ø)	⬆️
qucumber/rbm/binary_rbm.py	100% <100%> (ø)	⬆️
... and 17 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 035990e...25f7810. Read the comment docs.

[emerali]

The scope of this PR ended up being a lot more than just improving the DensityMatrix implementation and making Observables compatible with mixed states.

One major improvement (besides the two mentioned above) was properly vectorizing gradient rotations for batches of data. Prior to this PR, we were looping over every element of the batch to compute rotated gradients. In terms of computational efficiency, this was equivalent to doing SGD with a batch size of 1—we were still averaging the gradients over the batch before performing an update step, so the trajectory in parameter space should be the same before and after this update is applied.
This change allows users to actually get decent performance when running on GPU. We aren't quite at CPU levels of performance yet, however, since there's still the step where we move from GPU to CPU in order to make use of numpy's complex number support—we should add a cplx-torch version of prod just to see whether it'd work well enough on GPU.

Slightly less important improvements include:

• Optimizing code for training statistics: NLL and KL
• Fixing exact-NLL gradient implementations, hence fixing NLL gradient tests (NLL gradient tests now run by default)
• Add more tests for Observables, and tweak SWAP observable such that the output of apply has the same length as the input batch of samples
• Cleaning up/reducing code duplication, particularly in tests
• Changed some arg/kwarg names in training_statistics; added the relevant deprecation warnings so as not to break backwards compatibility.
• Re-arrange tutorials, Density Matrix tutorial is now 3rd, the observables tutorials were pushed back a little.