added references

README.md

+# Learning kinetic energy functional through differentiable optimization
+
+Here, the key is to include the integrated density difference in the loss function. Training this then require one to 
+differentiate through the Euler equation solver for the orbital-free density functionals which can be easily done with `jaxopt.implicit_diff.custom_root` of https://arxiv.org/abs/2105.15183. 
+
+Traning will be more expensive (10x slow than just use total energy and its functional derivative in the loss), but the hope is the trained kinetic energy functional is more generalizable. 
+
+## Requirements
+
+```bash
+pip install git+https://github.com/google/jaxopt
+pip install git+https://github.com/deepmind/dm-haiku
+```
+
+
+## To run
+
+```python
+python src/train.py --k 100 --b 32 
+python src/inference.py --restore_path  --k 100 --b 32 data/kejax_N1_G300_d2_c16_k100_b32_lr0.0001
+```

src/data.py

 import jax
 import jax.numpy as jnp
 
+'''
+cf https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.108.253002
+'''
 def make_training_data(N, G):
 
     xmesh = jnp.linspace(0, 1, G, endpoint=True)   

src/model.py

@@ -2,6 +2,9 @@ import jax
 import jax.numpy as jnp
 import haiku as hk 
 
+'''
+cf fig1b of https://pubs.acs.org/doi/10.1021/acs.jctc.0c00580
+'''
 def make_network(depth, channels, kernelsize):
 
     def forward_fn(x):