Google Summer of Code 2022: Ferminet Model implementation for deepchem

Hey everyone!

I am Shaipranesh S. I will be working with DeepChem this summer as a GSoC student.

About me

I am a sophomore at the Birla Institute of Technology and Science, Pilani. I am passionate about Machine Learning and am interested in building Machine Learning models for scientific applications. Since February, I have been a part of DeepChem’s community and am incredibly excited to work with all of you and hoping!

Contact Details

About the Project

My project aims to add a model to deepchem called the ferminet, a deep-learning-based variational Monte Carlo technique used to add accuracy to wavefunctions obtained from the Hartree-Fock method by including the electron correlation factor. The model is a permutation equivariant neural network and involves a second-order optimizer and an MCMC sampler.

Project description is available here.

I am excited about this project, and do follow this thread to see my weekly updates!!

4 Likes

General updates for this week:
-Updated 1 PR: Adding DeepQMC tutorial by shaipranesh2 · Pull Request #2914 · deepchem/deepchem (github.com)

Updated the description text and code according to the suggestions gave by Tony and Peter.

-Issued a new PR: Adding Metropolis Hasting sampler[WIP] by shaipranesh2 · Pull Request #2935 · deepchem/deepchem (github.com)

Put up a draft PR for the MCMC sampler for electron sampling.

Had a few questions regarding the asymmetric sampling of the sampler and got it clarified with Dr. David Pfau(one of the authors of Ferminet).

And have some ambiguity for sampling for ions, which I hope to get it cleared within next week on how the electrons lost or gained should be added to which atom of the molecule specifically? (This might not be an issue for diatomic and monoatomic ions)

Future work:

-Adding Unit tests for electron sampler and formatting to deepchem guidelines.

-Getting started on adding Second-order optimizers(K-FAC and SGD) to deepchem

Also, this week a fair bit of my time was spent on reading Metropolis-Hastings and related to concepts to understand it better. Also, Dr. Pfau mentioned about Hamiltonian Monte Carlo on how it can give improved results but had trouble as it did not give meaningful samples.(could look into it after the basic implementation for a better result)

Progress 11-17 June:
-Made all changes to the DeepQMC tutorial as discussed and got it merged in!
-Updated PR https://github.com/deepchem/deepchem/pull/2935:
-Fixed bugs in simultaneous MCMC move and is working fine(both symmetric and asymmetric)
-Added gauss initialize function(keeping in mind of the ions case discussed in the past week)
-Sampler now works for any distribution of 1-D array of any dimensions
-Shifted the sampler to utils
To-do:
- Debug single-electron move
-Adding unit tests(have already written tests for some functions that I used when debugging)
-Write documentation in utils.rst
-Start reading on K-FAC and should implement it

1 Like

Progress 18-24 June

-Updated PR https://github.com/deepchem/deepchem/pull/2935:

-The class gives outputs as expected

-Addressed the reviews and made it as an ElectronSampler class

-Did changes to API as discussed and added docstrings and text.

  • K-FAC:

-Did initial reading on second order optimizers and in particular K-FAC.

  • A repo( https://github.com/alecwangcq/KFAC-Pytorch) had K-FAC implemented, mailed the author about it.

  • A question on Functorch library on optimizing K-FAC

  • Saw Jax had features like automatic vectorization and is really fast at calculating hessians. Also, saw that ElectronSampler

had plenty of random normal initializations and summation within an axis of array ,where numpy was found to be faster than Jax with JIT in GPU, hence

I felt numpy to be useful in this case.

To-do:

-K-FAC implementation

1 Like