Roadmap

[PhilipVinc]

I would like to keep here a list of issues and enhancements to consider for the future

New Features:

• Add dynamics Add time evolution driver with jax-based ODE solver #1012 (Done)
• Add Quantum State Reconstruction Implement quantum state reconstruction (pure & mixed) #1427 (Done)
• Support gradient of non-hermitian operators (the code is already implemented, but there might be bugs and it's disabled)
• Implement gradient of expectation values of operators on mixed states
• Create a new operator, Covariance/Correlator, and its kernel to compute the expectation value quickly and with the right variance

Models:

• Add Symmetric Dense layer to have RBMSymm Add basic permutation-invariant RBMSymm implementation #555 (Done)
  • Make the interface to construct them simpler Add HashableArray utility class #621, Turn graph.SymmGroup into generic PermutationGroup #709
• Add Autoregressive FFNN Autoregressive model and sampler #705 (Done)

Hilbert:

• Find a better way to work with DoubledHilbert and to construct it. Make DoubledHilbert a parametric class #788 (Done)
• Allow a mechanism to override the random_state and flip_site for certain doubled hilbert, such as homogeneous ones. They don't need to go through the jax.lax.switch all the time Make DoubledHilbert a parametric class #788 (Done)
• Define HomogeneousHilbert abstract type and inherit from it [NFC] move custom_hilbert common logic to HomogeneousHilbert #735 (Done)

Samplers:

• Autodetect CPUs/GPUs and decide to which sampler to dispatch
• Use Numba in samplers Numba4jax: Numba in HamiltonianSampler #554 (Done)
  • Use Numba in GPU samplers remove nk.jax.numba4jax #934 (Done by switching to numba4jax)
  • Use Numba in Parallel Tempered samplers (Done)

Variational States:

• Make padding optional, give a way to control it.
• In general, find a way to work with operators where padding is not an on/off switch, but rather works in increments of an arbitrary quantity (say 512)...
• Consider jitting model construction jit model init #832 (Done)

Tests:

• Add benchmarks to avoid performance regressions Benchmarks? #68
• Eventually consider benchmarking also jit-compile time?
• Properly test Liouvillian gradient (done)

Misc

• Possibly improve jit time of nk.sampler.sample by refactoring the code and where we jit
  • Partially addressed in [Internal] Changes to (semi-)internal sampler api #721
• Keep an eye out for serialising jax jitted functions: Jax#476
• ~Upstream what we need for complex number support to Flax and eventually remove nk.nn ~ (Done)
  • In progress Complex-valued layers (e.g: dense) google/flax#805
  • It was decided that we will never remove nk.nn, but this will eventually only contain our own stuff, not vendored versions of dependencies.
• Some utils to deserialise our .log files and reconvert to our history objects and make plotting easier might be nice.
  • Partly done for deserialising, not for plotting.

Issues:

• MCState.samples before and after serialization/deserialization don't match. Right now, when serialising a variational state, we store the parameters and the sampler_state. However, if we have already sampled the state, the current sampler_state will generate the next batch of samples. So we should keep around the last sampler_state and serialise that one, so that when calling .samples after deserialisation we get the same samples. (Done)

Docs:

• Docs have plenty of build warnings. Those should be addressed (Done)
• Add Developer documentation
• Move the core of the docs (not the website) to ReadTheDocs for automatic versioning? (Done)

[gcarleo]

Yes, I would add

• Remove Legacy [Wait] Remove legacy #773