Make pretrained.py (or similar) usable for ALIGNN FF

[ironhammer269]

As of right now, the only way to perform inference on ALIGNN FF models is to use the run_alignn_ff script.

It is practical if one is interested in performing MD on a single structure, but for prediction of simple properties for multiple structures (eg. unrelaxed energy), it is very inefficient since one has to reload the model for every inference. Also, it is not easily parallelized as the code for calculating the energy is locked in another module after "self.atoms = self.jarvis_atoms.ase_converter()" line 367 in ff.py.

Thus i have been trying to use pretrained.py for an ALIGNN FF model but i am running into the following issue :

(alignn_env) PS C:\Users\haear\Documents\GitHub\alignn\alignn> python .\pretrained.py --model_path "C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed" --file_format poscar --file_path "C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\DFT\test_multi_pass\POSCAR1"
Using chk file C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed\best_model.pt
Path C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed
Config C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed\config.json
C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\dgl\backend\pytorch\tensor.py:445: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly. To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
assert input.numel() == input.storage().size(), (
Traceback (most recent call last):
File "C:\Users\haear\Documents\GitHub\alignn\alignn\pretrained.py", line 506, in
out_data = get_prediction(
File "C:\Users\haear\Documents\GitHub\alignn\alignn\pretrained.py", line 344, in get_prediction
model([g.to(device), lg.to(device)])['out']
File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\alignn\models\alignn_atomwise.py", line 524, in forward
/ g.ndata["V"][count_node + num_nodes]
File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\dgl\view.py", line 80, in getitem
return self._graph._get_n_repr(self._ntid, self._nodes)[key]
File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\dgl\frame.py", line 688, in getitem
return self._columns[name].data
KeyError: 'V'

I guess this is linked to the difference in construction of the graphs between ALIGNN FF/ non FF, but i haven't been able to track down what it is. Also this difference is mainly a decoration since both models should take the same input for inference (ie. a simple POSCAR file).

If anyone has insight on how to fix this, I would be very grateful.