Add sGNN-PF6 example

examples/sgnn/PF6.pdb

@@ -0,0 +1,19 @@
+HEADER    
+TITLE     MDANALYSIS FRAME 0: Created by PDBWriter
+REMARK     Created by AuToFF
+CRYST1   50.000   50.000   50.000  90.00  90.00  90.00 P 1           1
+ATOM      1  F01 PF6 1   1       1.069  -0.257   0.625  1.00  0.00      SYST F  
+ATOM      2  P02 PF6 1   1       2.483   0.187   0.143  1.00  0.00      SYST P  
+ATOM      3  F03 PF6 1   1       1.556   0.298  -1.104  1.00  0.00      SYST F  
+ATOM      4  F04 PF6 1   1       2.495   1.222   1.308  1.00  0.00      SYST F  
+ATOM      5  F05 PF6 1   1       3.427  -0.713  -0.710  1.00  0.00      SYST F  
+ATOM      6  F06 PF6 1   1       3.264   1.422  -0.397  1.00  0.00      SYST F  
+ATOM      7  F07 PF6 1   1       3.085  -0.854   1.134  1.00  0.00      SYST F  
+CONECT    1    2
+CONECT    2    1    3    4    5    6    7
+CONECT    3    2
+CONECT    4    2
+CONECT    5    2
+CONECT    6    2
+CONECT    7    2
+END

examples/sgnn/test_PF6.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+import time
+import sys
+import jax
+import jax.numpy as jnp
+from jax import grad, value_and_grad, jit
+import numpy as np
+import dmff
+from dmff.utils import jit_condition
+# from gnn import MolGNNForce
+from dmff.sgnn.gnn import MolGNNForce
+# from graph import TopGraph, from_pdb
+from dmff.sgnn.graph import TopGraph, from_pdb
+import optax
+import pickle
+# use pytorch data loader
+from torch.utils.data import DataLoader
+from jax.lib import xla_bridge 
+print(jax.devices()[0]) 
+print(xla_bridge.get_backend().platform)
+
+
+if __name__ == "__main__":
+    box = jnp.eye(3) * 50
+    pdbfile = 'PF6.pdb'
+    # pdbfile = 'pdb_bank/BF4.pdb'
+    # pdbfile = 'pdb_bank/DFP.pdb'
+    # pdbfile = 'pdb_bank/FSI.pdb' # no problem
+    G = from_pdb(pdbfile)
+    model = MolGNNForce(G, nn=0, atype_index={'P': 0, 'F': 1}, max_valence=6)
+    energy = model.forward(G.positions, box, model.params)
+    print(energy)