Google Summer of Code 2022: Ferminet Model implementation for deepchem

#1

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
#2

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)

#3

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
#4

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
#5

Progress 25th June-1st July

-Updated PR https://github.com/deepchem/deepchem/pull/2935 and looks fine to merge in.
-This will add Electron initialization & Sampling using Gauss distribution sampling and Markov Chain Monte Carlo, respectively.
-Created the new PR https://github.com/deepchem/deepchem/pull/2972.
-Added the KFAC implementation from https://github.com/alecwangcq/KFAC-Pytorch/tree/master/optimizers and modified a little according to ferminet’s needs(added mean centered covariance as an option).
-Added KFAC optimizer implementation.
-Added tests
To-do:
- Add docstrings
-Address some of the issues(like adding damping to the momentum)

1 Like
#6

Progress 1 - 8 July

-Updated https://github.com/deepchem/deepchem/pull/2972 : KFAC PR
-Added documentations, test and few bit of a modifications to the KFAC optimizer
-Did some changes to torch model class, so that when using KFAC, the model is given input instead of the model’s parameters
-Did some minor changes to the electron sampler utility class, which was merged at start of this week:
- https://github.com/deepchem/deepchem/pull/2987: PEP-008 style corrections and some changes to pass doctest
- https://github.com/deepchem/deepchem/pull/2973: utils.rst update
-Started on the Ferminet model implementation. Have written some tests, docstrings and the input featurizer(the one electron and two electron featurizers)

1 Like
#7

Progress 9 - 15 July

Updated https://github.com/deepchem/deepchem/pull/2972 : KFAC PR

  • Addressed reviews, fixed some bugs
  • Passed its tests - fitting on a dummy data using a CNN(tests both Linear and Conv2d layers)
  • Has some mypy issues.

Started on Ferminet model implementation, which I aim to complete via several small PRs:
-Added PR #2997 :https://github.com/deepchem/deepchem/pull/2997
- Aims to sample electrons and perform calculations to give single and double electron stream

#8

Progress July 24-29
https://github.com/deepchem/deepchem/pull/2972 : KFAC-PR:
-resolved the mypy errors using reveal_type
-Ready to merge in!!

  • https://github.com/deepchem/deepchem/pull/2997: Ferminet input stream PR
    • Gives both one-electron and two-electron featurizer from a list of atom’s SMILES and coordinates.
      -Added tests and docstring
      -Ready to merge in after a review
  • Working on the loss function of Ferminet(hamiltonian). Have the potential energy part ready, will put up a PR once input stream gets merged in.

TODO:
Complete the loss function by adding the kinetic energy operator too
Start working on Pyscf for pre-training

#9

Progress July 30 - Aug 5

  • https://github.com/deepchem/deepchem/pull/2997: Ferminet input stream PR
    -Did some changes - added spin-wise arrays
    -passed all the tests
    -Still have some ambiguity over the ionic charge in the molecule system
    -Wrote code for calculating potential energy of the system, will merge after input stream gets finalized
    -Todo
    -Complete Hamiltonian loss function by including the KE too.
1 Like
#10

Progress Aug 6 - Aug 12

-Added support for ionic charges

1 Like
#11

Progress Aug 12 - Aug 19

TODO:- Work on the pretraining(Can be used to test the model as well)

#12

Progress Aug 19 - Aug 26

-Have some more mypy issues

-some doubts on model fitting

1 Like
#13

Progress Aug 26 - Sep 2

Ferminet loss and model prototype PR

  • Debugging dimension issues
  • Rewrote part of the numpy code into torch

Todo:-

-Batching

-Pretraining

#14

Progress Sep 2 - 9

Ferminet loss and model prototype PR

  • Did changes to the layers, now the model.forward works for 1 iteration, have to incorporate changes for further iterations
  • Looked into Pyscf for getting Hartree-Fock baseline, and should make changes to initial sampling for pretraining.
#15

Progress Sep 9 - 16

Ferminet loss and model prototype PR

  • All tests passed
  • model.forward() is running without errors
  • Gives log probability of square of wave-function as the output
  • Ready to merge in?
1 Like
#16

Progress Sep 16 - 23

Ferminet loss and model prototype PR

  • Overriden fit method from torch_model
  • Batching done
  • Loss function complete, but getting NaNs for jacobian.

Pretraining with PySCF would be the last step.

1 Like