GSoC 2020 : JAX Chem project

Hello, my name is Daiki Nishikawa, please call me Daiki!
I’m excited to be selected as a member of GSoC this year and looking forward to work with deepchem community members!

About me

I’m a second year master’s student and belong to the cheminformatics laboratory at the University of Tokyo. My current research is to design new inorganic materials using virtual screening with graph convolutional networks or deep generative models. My bachelor course research was the analysis of the electrochemical reaction by using a quantum chemical simulation. (I used VASP and Gaussian.)

Strong area

  • Modeling graph convolutional networks
    • I worked for Chainer Chemistry at Preferred Networks last summer
  • Treating material data
    • I get used to Materials Project Dataset and Pymatgen
  • OSS activity
    • I sent some PRs to various OSS, ex) React Native

Current Interest

  • Deep generative models, especially disentangled VAE
  • Deep metric learning
  • OSS activity
  • MLOps

Links

About the GSoC project

In this GSoC, I will launch the library called JAX-Chem, which will make it easy to do various chemical and physical modelings with JAX. My motivation for JAX-Chem is to provide an alternative option when implementing deep learning models about chemistry. JAX is developed by Google and most APIs are like Numpy. So, I think it is easy for many people to get used to JAX and read complex modeling codes compared with TensorFlow and PyTorch.

During GSoC, I will be developing in https://github.com/deepchem/jaxchem repository. (not in the main deepchem repo) As a first step, I will implement the GCN benchmark and launch this project because I understood the GCN deeply during last summer internship. After that, I will also try to implement other models like Attention-based models or related to physical phenomenons (if possible, I want to pick up the surface reaction modeling).

Current Implementation plan in JAX-chem

  • GCN benchmarks
    • Gated Graph Neural Network
    • Crystal Graph Convolutional Neural Network
  • Attention-based models
    • Molecular Attention Transformers

More details

2 Likes

Hello Daiki
I am looking forward to see your contributions.
I am also interested in implementing chemical models in JAX. I have started learning JAX a month ago, and hopefully, I am going to implement some of the Deepchem models (and other models beyond Deepchem) on JAX.

3 Likes

@HFooladi

Now, we welcome to contribute the models on JAX.
Please see the JAXChem repository!

2 Likes