Pure Python example for prediction

[mhellstr]

It would be helpful/instructive if there was a pure Python example that does not rely on LAMMPS for prediction.

For example, how could I use benchmark.pt to get energy/forces/stress of a single structure if I have numpy arrays with atomic numbers, coordinates, lattice vectors?

[muzhenliang]

You can load model by:

m = torch.jit.load("path/to/model.pt")

data = {"input": "your_input_data[in format torch tensor]"}

# Required keys:
# edge_index: [2, Ne].
# edge_cell_shift: [Ne, 3]. Unit: Angstrom
# coul_edge_index: [2, Ne_coul]. 
# coul_edge_cell_shift: [Ne_coul, 3]. Unit: Angstrom
# disp_edge_index: [2, Ne_disp]. 
# disp_edge_cell_shift: [Ne_disp, 3]. Unit: Angstrom
# atom_types: [Na]. torch.long
# cell: [3, 3]. Unit: Angstrom. None or negative means non-PBC
# g_ewald: [1]. g_ewald parameter in LAMMPS, could be 1.0 for test

pred = m.forward(data)

# pred is the output of the model```

[mhellstr]

Thanks, but I guess the question is how would I know how to populate edge_index, edge_cell_shift, etc. ?