https://github.com/tummfm/jax-dimenet

Jax / Haiku implementation of DimeNet++.

https://github.com/tummfm/jax-dimenet

Science Score: 23.0%

This score indicates how likely this project is to be science-related based on various indicators:

  • CITATION.cff file
  • codemeta.json file
  • .zenodo.json file
  • DOI references
    Found 1 DOI reference(s) in README
  • Academic publication links
  • Committers with academic emails
    1 of 1 committers (100.0%) from academic institutions
  • Institutional organization owner
  • JOSS paper metadata
  • Scientific vocabulary similarity
    Low similarity (14.8%) to scientific vocabulary

Keywords

graph-neural-networks jax molecular-dynamics
Last synced: 10 months ago · JSON representation

Repository

Jax / Haiku implementation of DimeNet++.

Basic Info
Statistics
  • Stars: 18
  • Watchers: 1
  • Forks: 1
  • Open Issues: 1
  • Releases: 0
Topics
graph-neural-networks jax molecular-dynamics
Created over 4 years ago · Last pushed over 4 years ago
Metadata Files
Readme License

README.md

Jax-DimeNet++

PyPI version Documentation Status

Haiku implementation of the DimeNet++ architecture.

Getting started

This repository provides 2 interfaces for the DimeNet++ model. The first interface allows using DimeNet++ as a Jax M.D. energy_fn to run molecular dynamics simulations. The second interface allows prediction of global molecular properties.

```python from jaxdimenet import dimenet, sparsegraph

Jax M.D. energy function:

initfn, dimenetenergyfn = dimenet.energyneighborlist(displacementfn, rcut) initparams = initfn(random.PRNGKey(0), positions, neighbor=neighbors) energyfn = partial(dimenetenergyfn, initparams) # jaxmd energyfn interface print('Predicted energy:', energy_fn(positions, neighbors))

Molecular property prediction:

molgraph, _ = sparsegraph.sparsegraphfromneighborlist( displacementfn, positions, neighbors, rcut) initfn, propertypredictor = dimenet.propertyprediction(rcut, ntargets=5) initparams = initfn(random.PRNGKey(0), molgraph) print('Predicted properties:', propertypredictor(initparams, molgraph)) ```

A minimum usage example is available in the notebooks folder. For more real-world applications of the DimeNet++ model in MD simulations, please refer to the DiffTRe repository.

Installation

You can install Jax-DimeNet via pip: pip install jax-dimenet

Requirements

The repository uses the following packages: jax>=0.2.12 jaxlib>=0.1.65 jax-md>=0.1.13 dm-haiku>=0.0.4 sympy chex The code was run with Python 3.8.

Contribution

Contributions are always welcome! Please open a pull request to discuss the code additions.

Contact

For questions, please contact stephan.thaler@tum.de.

Citation

If you use this code in your own work, please consider citing the paper that introduced this DimeNet++ implementation, text @article{thaler_difftre_2021, title = {Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting}, author = {Thaler, Stephan and Zavadlav, Julija}, journal={Nature Communications}, volume={12}, pages={6884}, doi={10.1038/s41467-021-27241-4} year = {2021} } as well as the original DimeNet++ paper. text @inproceedings{klicpera_dimenetpp_2020, title = {Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules}, author = {Klicpera, Johannes and Giri, Shankari and Margraf, Johannes T. and G{\"u}nnemann, Stephan}, booktitle={NeurIPS-W}, year = {2020} }

Owner

  • Name: Multiscale Modeling of Fluid Materials
  • Login: tummfm
  • Kind: organization
  • Location: Garching near Munich, Germany

Technical University of Munich TUM

GitHub Events

Total
  • Watch event: 2
  • Issue comment event: 1
Last Year
  • Watch event: 2
  • Issue comment event: 1

Committers

Last synced: over 3 years ago

All Time
  • Total Commits: 6
  • Total Committers: 1
  • Avg Commits per committer: 6.0
  • Development Distribution Score (DDS): 0.0
Top Committers
Name Email Commits
stephan s****r@t****e 6
Committer Domains (Top 20 + Academic)
tum.de: 1

Issues and Pull Requests

Last synced: 10 months ago

All Time
  • Total issues: 1
  • Total pull requests: 0
  • Average time to close issues: N/A
  • Average time to close pull requests: N/A
  • Total issue authors: 1
  • Total pull request authors: 0
  • Average comments per issue: 1.0
  • Average comments per pull request: 0
  • Merged pull requests: 0
  • Bot issues: 0
  • Bot pull requests: 0
Past Year
  • Issues: 1
  • Pull requests: 0
  • Average time to close issues: N/A
  • Average time to close pull requests: N/A
  • Issue authors: 1
  • Pull request authors: 0
  • Average comments per issue: 1.0
  • Average comments per pull request: 0
  • Merged pull requests: 0
  • Bot issues: 0
  • Bot pull requests: 0
Top Authors
Issue Authors
Pull Request Authors
Top Labels
Issue Labels
Pull Request Labels

Packages

  • Total packages: 1
  • Total downloads:
    • pypi 11 last-month
  • Total dependent packages: 0
  • Total dependent repositories: 1
  • Total versions: 1
  • Total maintainers: 1
pypi.org: jax-dimenet

DimeNet++ in Jax.

  • Versions: 1
  • Dependent Packages: 0
  • Dependent Repositories: 1
  • Downloads: 11 Last month
Rankings
Dependent packages count: 10.1%
Stargazers count: 15.3%
Dependent repos count: 21.6%
Forks count: 22.7%
Average: 28.6%
Downloads: 73.6%
Maintainers (1)
Last synced: 10 months ago

Dependencies

docs/requirements.txt pypi
  • myst_nb *
  • sphinx >=3
  • sphinx-autodoc-typehints *
  • sphinx_rtd_theme *
setup.py pypi
  • chex *
  • dm-haiku >=0.0.4
  • jax >=0.2.12
  • jax-md >=0.1.13
  • jaxlib >=0.1.65
  • sympy *