Skip to content

Roadmap

Created by: PhilipVinc

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

New Features:

  • Add dynamics #1012 (Done)
  • Add Quantum State Reconstruction #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 #555 (Done)
    • Make the interface to construct them simpler #621, #709
  • Add Autoregressive FFNN #705 (Done)

Hilbert:

  • Find a better way to work with DoubledHilbert and to construct it. #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 #788 (Done)
  • Define HomogeneousHilbert abstract type and inherit from it #735

Samplers:

  • Autodetect CPUs/GPUs and decide to which sampler to dispatch
  • Use Numba in samplers #554 (Done)
    • Use Numba in GPU samplers #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 #832 (Done)

Tests:

  • Add benchmarks to avoid performance regressions #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 #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
    • In progress 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.

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
  • Add Developer documentation
  • Maybe move the core of the docs (not the website) to ReadTheDocs for automatic versioning?