chemtrain.neural_networks.rst

Neural Networks

chemtrain.neural_networks

DimeNet++

The DimeNetPP directly takes a ~chemtrain.sparse_graph.SparseDirectionalGraph as input and predicts per-atom quantities. As DimeNetPP is a haiku Module, it needs to be wrapped inside a hk.transform() before it can be applied.

We provide 2 interfaces to DimeNet++: The function dimenetpp_neighborlist serves as a interface to Jax M.D. The resulting apply function can be directly used as a jax_md energy_fn, e.g. to run molecular dynamics simulations.

For direct prediction of global molecular properties, dimenetpp_property_prediction can be used.

DimeNetPP

__init__

__call__

dimenetpp_neighborlist

dimenetpp_property_prediction

Pairwise NN

PairwiseNN implements a neural network, that parametrizes 2-body interactions. The function pair_interaction_nn initializes a pairwise jax_md neighborlist energy_fn, as an alternative to classical tabulated potentials.

PairwiseNN

__init__

__call__

pair_interaction_nn