Fix file format.

alignn/tests/test_prop.py

@@ -9,8 +9,7 @@
 from sklearn.metrics import mean_absolute_error
 import os
 from jarvis.core.atoms import Atoms
-from alignn.train_folder import train_for_folder
-from alignn.train_folder_ff import train_for_folder as train_for_folder_ff
+from alignn.train_alignn import train_for_folder
 from jarvis.db.figshare import get_jid_data
 from alignn.ff.ff import AlignnAtomwiseCalculator, default_path, revised_path
 
@@ -175,7 +174,7 @@ def test_alignn_train():
             "../examples/sample_data_ff/config_example_atomwise.json",
         )
     )
-    train_for_folder_ff(root_dir=root_dir, config_name=config)
+    train_for_folder(root_dir=root_dir, config_name=config)
 
 
 def test_calculator():

alignn/train_alignn.py

@@ -215,7 +215,21 @@ def train_for_folder(
                 n_outputs.append(tmp)
             info["target"] = tmp
             file_path = os.path.join(root_dir, file_name)
-            atoms = Atoms.from_poscar(file_path)
+            if file_format == "poscar":
+                atoms = Atoms.from_poscar(file_path)
+            elif file_format == "cif":
+                atoms = Atoms.from_cif(file_path)
+            elif file_format == "xyz":
+                atoms = Atoms.from_xyz(file_path, box_size=500)
+            elif file_format == "pdb":
+                # Note using 500 angstrom as box size
+                # Recommended install pytraj
+                # conda install -c ambermd pytraj
+                atoms = Atoms.from_pdb(file_path, max_lat=500)
+            else:
+                raise NotImplementedError(
+                    "File format not implemented", file_format
+                )
             info["atoms"] = atoms.to_dict()
         else:
             info["target"] = i[target_key]