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 (Done)
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 ~ (Done)
~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.
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)
I would like to keep here a list of issues and enhancements to consider for the future
New Features:
Models:
Hilbert:
DoubledHilbert
and to construct it.~ #788 (Done)random_state
andflip_site
for certain doubled hilbert, such as homogeneous ones. They don't need to go through thejax.lax.switch
all the time~ #788 (Done)Samplers:
Variational States:
Tests:
Misc
nk.sampler.sample
by refactoring the code and where we jitnk.nn
~ (Done)nk.nn
, but this will eventually only contain our own stuff, not vendored versions of dependencies..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 thesampler_state
. However, if we have already sampled the state, the currentsampler_state
will generate the next batch of samples. So we should keep around the lastsampler_state
and serialise that one, so that when calling.samples
after deserialisation we get the same samples.~ (Done)Docs: