New trainer compatibility fixes

TRAIN.md

@@ -204,11 +204,11 @@ To train and validate an OC20 IS2RE/S2EF model on total energies instead of adso
 
 ```yaml
 task:
-  dataset: oc22_lmdb
   prediction_dtype: float32
   ...
 
 dataset:
+  format: oc22_lmdb
   train:
     src: data/oc20/s2ef/train
     normalize_labels: False
@@ -308,8 +308,8 @@ For the IS2RE-Total task, the model takes the initial structure as input and pre
 ```yaml
 trainer: energy # Use the EnergyTrainer
 
-task:
-  dataset: oc22_lmdb # Use the OC22LmdbDataset
+dataset:
+  format: oc22_lmdb # Use the OC22LmdbDataset
   ...
 ```
 You can find examples configuration files in [`configs/oc22/is2re`](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/oc22/is2re).
@@ -321,8 +321,8 @@ The S2EF-Total task takes a structure and predicts the total DFT energy and per-
 ```yaml
 trainer: forces  # Use the ForcesTrainer
 
-task:
-  dataset: oc22_lmdb # Use the OC22LmdbDataset
+dataset:
+  format: oc22_lmdb # Use the OC22LmdbDataset
   ...
 ```
 You can find examples configuration files in [`configs/oc22/s2ef`](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/oc22/s2ef).
@@ -332,8 +332,8 @@ You can find examples configuration files in [`configs/oc22/s2ef`](https://githu
 Training on OC20 total energies whether independently or jointly with OC22 requires a path to the `oc20_ref` (download link provided below) to be specified in the configuration file. These are necessary to convert OC20 adsorption energies into their corresponding total energies. The following changes in the configuration file capture these changes:
 
 ```yaml
-task:
-  dataset: oc22_lmdb
+dataset:
+  format: oc22_lmdb
   ...
 
 dataset:
@@ -382,10 +382,8 @@ If your data is already in an [ASE Database](https://databases.fysik.dtu.dk/ase/
 To use this dataset, we will just have to change our config files to use the ASE DB Dataset rather than the LMDB Dataset:
 
 ```yaml
-task:
-  dataset: ase_db
-
 dataset:
+  format: ase_db
   train:
     src: # The path/address to your ASE DB
     connect_args:
@@ -420,10 +418,8 @@ It is possible to train/predict directly on ASE-readable files. This is only rec
 This dataset assumes a single structure will be obtained from each file:
 
 ```yaml
-task:
-  dataset: ase_read
-
 dataset:
+  format: ase_read
   train:
     src: # The folder that contains ASE-readable files
     pattern: # Pattern matching each file you want to read (e.g. "*/POSCAR"). Search recursively with two wildcards: "**/*.cif".
@@ -443,10 +439,8 @@ dataset:
 This dataset supports reading files that each contain multiple structure (for example, an ASE .traj file). Using an index file, which tells the dataset how many structures each file contains, is recommended. Otherwise, the dataset is forced to load every file at startup and count the number of structures!
 
 ```yaml
-task:
-  dataset: ase_read_multi
-
 dataset:
+  format: ase_read_multi
   train:
     index_file: Filepath to an index file which contains each filename and the number of structures in each file. e.g.:
             /path/to/relaxation1.traj 200

ocpmodels/common/utils.py

@@ -1204,13 +1204,22 @@ def update_config(base_config):
     are now. Update old configs to fit the new expected structure.
     """
     config = copy.deepcopy(base_config)
-    config["dataset"]["format"] = config["task"].get("dataset", "lmdb")
+
+    # If config["dataset"]["format"] is missing, get it from the task (legacy location).
+    # If it is not there either, default to LMDB.
+    config["dataset"]["format"] = config["dataset"].get(
+        "format", config["task"].get("dataset", "lmdb")
+    )
+
     ### Read task based off config structure, similar to OCPCalculator.
     if config["task"]["dataset"] in [
         "trajectory_lmdb",
         "lmdb",
         "trajectory_lmdb_v2",
         "oc22_lmdb",
+        "ase_read",
+        "ase_read_multi",
+        "ase_db",
     ]:
         task = "s2ef"
     elif config["task"]["dataset"] == "single_point_lmdb":

ocpmodels/preprocessing/atoms_to_graphs.py

@@ -193,10 +193,10 @@ def convert(self, atoms: ase.Atoms, sid=None):
             data.cell_offsets = cell_offsets
         if self.r_energy:
             energy = atoms.get_potential_energy(apply_constraint=False)
-            data.y = energy
+            data.energy = energy
         if self.r_forces:
             forces = torch.Tensor(atoms.get_forces(apply_constraint=False))
-            data.force = forces
+            data.forces = forces
         if self.r_stress:
             stress = torch.Tensor(
                 atoms.get_stress(apply_constraint=False, voigt=False)

ocpmodels/trainers/ocp_trainer.py

@@ -529,7 +529,16 @@ def predict(
             axis = 0 if isinstance(predictions[key][0], np.ndarray) else None
             predictions[key] = np.concatenate(predictions[key], axis=axis)
 
-        self.save_results(predictions, results_file)
+        if self.is_debug:
+            try:
+                self.save_results(predictions, results_file)
+            except FileNotFoundError:
+                logging.warning(
+                    "Predictions npz file not found. "
+                    + "This file was not written since the trainer is running in debug mode."
+                )
+        else:
+            self.save_results(predictions, results_file)
 
         if self.ema:
             self.ema.restore()
@@ -582,7 +591,7 @@ def run_relaxations(self, split="val"):
             if check_traj_files(
                 batch, self.config["task"]["relax_opt"].get("traj_dir", None)
             ):
-                logging.info(f"Skipping batch: {batch[0].sid.tolist()}")
+                logging.info(f"Skipping batch: {list(batch[0].sid)}")
                 continue
 
             relaxed_batch = ml_relax(
@@ -597,7 +606,7 @@ def run_relaxations(self, split="val"):
             )
 
             if self.config["task"].get("write_pos", False):
-                systemids = [str(i) for i in relaxed_batch.sid.tolist()]
+                systemids = [str(i) for i in list(relaxed_batch.sid)]
                 natoms = relaxed_batch.natoms.tolist()
                 positions = torch.split(relaxed_batch.pos, natoms)
                 batch_relaxed_positions = [pos.tolist() for pos in positions]

tests/datasets/test_ase_datasets.py

@@ -117,7 +117,7 @@ def test_ase_dataset(ase_dataset):
     assert len(dataset) == mult * len(structures)
     for data in dataset:
         assert hasattr(data, "y")
-        assert data.force.shape == (data.natoms, 3)
+        assert data.forces.shape == (data.natoms, 3)
         assert data.stress.shape == (3, 3)
         assert data.tensor_property.shape == (6, 6)
         assert isinstance(data.extensive_property, int)
@@ -249,8 +249,8 @@ def test_ase_multiread_dataset(tmp_path) -> None:
     assert len(dataset) == len(atoms_objects)
 
     assert hasattr(dataset[0], "y_relaxed")
-    assert dataset[0].y_relaxed != dataset[0].y
-    assert dataset[-1].y_relaxed == dataset[-1].y
+    assert dataset[0].y_relaxed != dataset[0].energy
+    assert dataset[-1].y_relaxed == dataset[-1].energy
 
     dataset = AseReadDataset(
         config={
@@ -268,4 +268,4 @@ def test_ase_multiread_dataset(tmp_path) -> None:
     )
 
     assert hasattr(dataset[0], "y_relaxed")
-    assert dataset[0].y_relaxed != dataset[0].y
+    assert dataset[0].y_relaxed != dataset[0].energy

tests/preprocessing/test_atoms_to_graphs.py

@@ -113,11 +113,11 @@ def test_convert(self) -> None:
         np.testing.assert_allclose(act_positions, positions)
         # check energy value
         act_energy = self.atoms.get_potential_energy(apply_constraint=False)
-        test_energy = data.y
+        test_energy = data.energy
         np.testing.assert_equal(act_energy, test_energy)
         # forces
         act_forces = self.atoms.get_forces(apply_constraint=False)
-        forces = data.force.numpy()
+        forces = data.forces.numpy()
         np.testing.assert_allclose(act_forces, forces)
         # stress
         act_stress = self.atoms.get_stress(apply_constraint=False, voigt=False)
@@ -145,11 +145,11 @@ def test_convert_all(self) -> None:
         np.testing.assert_allclose(act_positions, positions)
         # check energy value
         act_energy = self.atoms.get_potential_energy(apply_constraint=False)
-        test_energy = data_list[0].y
+        test_energy = data_list[0].energy
         np.testing.assert_equal(act_energy, test_energy)
         # forces
         act_forces = self.atoms.get_forces(apply_constraint=False)
-        forces = data_list[0].force.numpy()
+        forces = data_list[0].forces.numpy()
         np.testing.assert_allclose(act_forces, forces)
         # stress
         act_stress = self.atoms.get_stress(apply_constraint=False, voigt=False)