coarse-graining / cgnet

learning coarse-grained force fields
BSD 3-Clause "New" or "Revised" License
53 stars 26 forks source link

Schnet layer #34

Closed nec4 closed 4 years ago

nec4 commented 5 years ago

see https://aip.scitation.org/doi/10.1063/1.5019779

Schnet is a variant of deep tensor neural networks (DTNNs). The main characteristic of Schnet is its interaction block which involves the following operations: schnet_arch

The interaction block is characterized by a residual branch:

  1. Dense layer
  2. Continuous filter convolution + residual connection of cartesian coordinates (r1...rn) forwarded through filter generating network
  3. Dense layer
  4. Shifted Softplus activation
  5. Dense layer

The output of this residual branch is added back to the input features (x1...xn) from the previous block, which is the forwarded through to the next block of the network. Essentially, there are two residual connections to account for: the first connects back to the cartesian coordinates through the filter generator, while the second connects to the feature inputs from the previous block.

Regarding the filter generator, it is essentially used in place of a filter-tensor (much like those found in convolutional networks). More accurately, the filter generating network is a special case of a factorized tensor layer found in traditional DTNN architectures.

nec4 commented 5 years ago

As for the structure of the filter generating network, here is the diagram from the paper: filter_gen

This block is characterized by the following sequential operations:

  1. Featurization of cartesian positions into distances (which are roto-translationally invariant)
  2. Expansion into a set of radial basis functions
  3. Dense layer
  4. Shifted softplus activation
  5. Dense layer
  6. Shifted softplus activation
  7. Optional periodic boundary condition pooling (used for systems on a lattice)
nec4 commented 5 years ago

I am thinking:

  1. FilterGeneratorNetwork class in nnet.py
  2. RBFlayer class in layers.py
  3. ShiftedSoftplus class in layers.py
  4. CFilterConv class in layers.py`

maybe a class for wrapping the interaction block too. Or perhaps is easier to adapt from the Schnet code? Not sure if they have a repo somewhere.

nec4 commented 5 years ago

See #35.

Dom1L commented 5 years ago

I made a PR with some pseudo-code to outline the SchNet implementation.

See #42

Dom1L commented 4 years ago

I think now that most of the tools are already implemented via #47 and #53 we don't need this issue anymore.

Further discussions about the implementation/integration are continued in issue #80