Use str ids in AseReadDataset (#1555)

* use str as system ids in AseReadDataset

* add AseReadDataset to tests

* reduce number of relax steps in test

* mark tests GPU

* fix tests

Author: lbluque

ruff.toml

@@ -33,6 +33,7 @@ select = [
   "YTT",    # flake8-2020
 ]
 ignore = [
+  "C408",    # Unnecessary `dict` call (rewrite as a literal)
   "PLR",    # Design related pylint codes
   "E501",   # Line too long
   "B028",   # No explicit stacklevel
@@ -55,7 +56,7 @@ ignore = [
   "RUF005", # concat lists
   "SIM108", # Use ternary operator
   "PT006", # Wrong type passed to first argument
-  "PYI024", # Use `typing.NamedTuple` instead of `collections.namedtuple`
+  "PYI024", # Use `typing.NamedTuple` instead of `collections.namedtuple`,
 ]
 unfixable = [
   "T20",  # Removes print statements

src/fairchem/core/datasets/ase_datasets.py

@@ -235,14 +235,14 @@ def _load_dataset_get_ids(self, config) -> list[Path]:
 
     def get_atoms(self, idx: int) -> ase.Atoms:
         try:
-            str_file = self.ids[idx]
-            atoms = ase.io.read(str_file, **self.ase_read_args)
+            file_path = self.ids[idx]
+            atoms = ase.io.read(file_path, **self.ase_read_args)
         except Exception as err:
             warnings.warn(f"{err} occured for: {idx}", stacklevel=2)
             raise err
 
         if "sid" not in atoms.info:
-            atoms.info["sid"] = str_file
+            atoms.info["sid"] = str(file_path)
 
         return atoms
 

tests/core/components/conftest.py

@@ -7,20 +7,14 @@
 
 from __future__ import annotations
 
-from itertools import product
 from random import choice
 
-import numpy as np
 import pytest
-from ase.db import connect
-from pymatgen.core import Structure
-from pymatgen.core.periodic_table import Element
 
 from fairchem.core import FAIRChemCalculator, pretrained_mlip
-from fairchem.core.datasets import AseDBDataset
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def calculator() -> FAIRChemCalculator:
     uma_sm_models = [
         model for model in pretrained_mlip.available_models if "uma-s" in model

tests/core/components/test_calculate_runners.py

@@ -9,6 +9,8 @@
 
 import os
 
+import pytest
+
 from fairchem.core.components.calculate import (
     ElasticityRunner,
     RelaxationRunner,
@@ -17,6 +19,7 @@
 from fairchem.core.datasets.atoms_sequence import AtomsDatasetSequence
 
 
+@pytest.mark.gpu()
 def test_elasticity_runner(calculator, dummy_binary_dataset, tmp_path):
     elastic_runner = ElasticityRunner(
         calculator, input_data=AtomsDatasetSequence(dummy_binary_dataset)
@@ -55,6 +58,7 @@ def test_elasticity_runner(calculator, dummy_binary_dataset, tmp_path):
     assert len(results) == len(dummy_binary_dataset) // 2
 
 
+@pytest.mark.gpu()
 def test_singlepoint_runner(calculator, dummy_binary_dataset, tmp_path):
     # Test basic instantiation
     singlepoint_runner = SinglePointRunner(
@@ -75,7 +79,6 @@ def test_singlepoint_runner(calculator, dummy_binary_dataset, tmp_path):
         input_data=AtomsDatasetSequence(dummy_binary_dataset),
         calculate_properties=["energy", "forces"],
         normalize_properties_by={"energy": "natoms"},
-        save_target_properties=["energy"],
     )
     results_custom = singlepoint_runner_custom.calculate()
     assert len(results_custom) == len(dummy_binary_dataset)
@@ -95,6 +98,7 @@ def test_singlepoint_runner(calculator, dummy_binary_dataset, tmp_path):
     assert singlepoint_runner.save_state("dummy_checkpoint") is True
 
 
+@pytest.mark.gpu()
 def test_relaxation_runner(calculator, dummy_binary_dataset, tmp_path):
     # Test basic instantiation
     relaxation_runner = RelaxationRunner(
@@ -121,9 +125,8 @@ def test_relaxation_runner(calculator, dummy_binary_dataset, tmp_path):
         calculate_properties=["energy", "forces"],
         save_relaxed_atoms=False,
         normalize_properties_by={"energy": "natoms"},
-        save_target_properties=["energy"],
         fmax=0.1,  # relax_kwargs
-        steps=50,  # relax_kwargs
+        steps=5,  # relax_kwargs
     )
     results_custom = relaxation_runner_custom.calculate()
     assert len(results_custom) == len(dummy_binary_dataset)

tests/core/conftest.py

@@ -16,12 +16,14 @@
 import numpy as np
 import pytest
 import torch
+from ase.calculators.singlepoint import SinglePointCalculator
 from ase.db import connect
+from ase.io import write
 from pymatgen.core import Structure
 from pymatgen.core.periodic_table import Element
 from syrupy.extensions.amber import AmberSnapshotExtension
 
-from fairchem.core.datasets import AseDBDataset
+from fairchem.core.datasets.ase_datasets import AseDBDataset, AseReadDataset
 from fairchem.core.units.mlip_unit.mlip_unit import (
     UNIT_INFERENCE_CHECKPOINT,
     UNIT_RESUME_CONFIG,
@@ -192,27 +194,36 @@ def dummy_binary_dataset_path(tmpdir_factory, dummy_element_refs):
                 + 0.05 * rng.random() * dummy_element_refs.mean()
             )
             atoms = structure.to_ase_atoms()
-            db.write(
+            atoms.calc = SinglePointCalculator(
                 atoms,
-                data={
-                    "sid": f"structure_{i}",
-                    "energy": energy,
-                    "forces": rng.random((2, 3)),
-                    "stress": rng.random((3, 3)),
-                },
+                energy=energy,
+                forces=rng.random((2, 3)),
+                stress=rng.random((3, 3)),
             )
+            # write to the lmdb file
+            db.write(atoms, data={"sid": f"structure_{i}"})
 
-    return tmpdir / "dummy.aselmdb"
+            # write it as a cif file as well
+            write(str(tmpdir / f"structure_{i}.cif"), atoms)
+
+    return tmpdir
+
+
+@pytest.fixture(scope="session", params=["asedb", "cif"])
+def dummy_binary_dataset(dummy_binary_dataset_path, request):
+    config = dict(src=str(dummy_binary_dataset_path))
+
+    if request.param == "cif":
+        config["pattern"] = "*.cif"
+        return AseReadDataset(config=config)
+    else:
+        return AseDBDataset(config=config)
 
 
 @pytest.fixture(scope="session")
-def dummy_binary_dataset(dummy_binary_dataset_path):
-    return AseDBDataset(
-        config={
-            "src": str(dummy_binary_dataset_path),
-            "a2g_args": {"r_data_keys": ["energy", "forces", "stress"]},
-        }
-    )
+def dummy_binary_db_dataset(dummy_binary_dataset_path):
+    config = dict(src=str(dummy_binary_dataset_path))
+    return AseDBDataset(config=config)
 
 
 @pytest.fixture(autouse=True)

tests/core/modules/test_element_references.py

@@ -26,10 +26,10 @@
 
 
 @pytest.fixture(scope="session", params=(True, False))
-def element_refs(dummy_binary_dataset, max_num_elements, request):
+def element_refs(dummy_binary_db_dataset, max_num_elements, request):
     return fit_linear_references(
         ["energy"],
-        dataset=dummy_binary_dataset,
+        dataset=dummy_binary_db_dataset,
         batch_size=16,
         shuffle=False,
         max_num_elements=max_num_elements,
@@ -39,12 +39,12 @@ def element_refs(dummy_binary_dataset, max_num_elements, request):
 
 
 def test_apply_linear_references(
-    element_refs, dummy_binary_dataset, dummy_element_refs
+    element_refs, dummy_binary_db_dataset, dummy_element_refs
 ):
     max_noise = 0.05 * dummy_element_refs.mean()
 
     # check that removing element refs keeps only values within max noise
-    batch = data_list_collater(list(dummy_binary_dataset), otf_graph=True)
+    batch = data_list_collater(list(dummy_binary_db_dataset), otf_graph=True)
     energy = batch.energy.clone().view(len(batch), -1)
     deref_energy = element_refs["energy"].dereference(energy, batch)
     assert all(deref_energy <= max_noise)
@@ -96,14 +96,14 @@ def test_create_element_references(element_refs, tmp_path):
 
 
 def test_fit_linear_references(
-    element_refs, dummy_binary_dataset, max_num_elements, dummy_element_refs
+    element_refs, dummy_binary_db_dataset, max_num_elements, dummy_element_refs
 ):
     # create the composition matrix
-    energy = np.array([d.energy for d in dummy_binary_dataset]).reshape(-1)
+    energy = np.array([d.energy for d in dummy_binary_db_dataset]).reshape(-1)
     cmatrix = np.vstack(
         [
             np.bincount(d.atomic_numbers.int().numpy(), minlength=max_num_elements + 1)
-            for d in dummy_binary_dataset
+            for d in dummy_binary_db_dataset
         ]
     )
     mask = cmatrix.sum(axis=0) != 0.0
@@ -130,30 +130,30 @@ def test_fit_linear_references(
     )
 
 
-def test_fit_seed_no_seed(dummy_binary_dataset, max_num_elements):
+def test_fit_seed_no_seed(dummy_binary_db_dataset, max_num_elements):
     refs_seed = fit_linear_references(
         ["energy"],
-        dataset=dummy_binary_dataset,
+        dataset=dummy_binary_db_dataset,
         batch_size=16,
-        num_batches=len(dummy_binary_dataset) // 16 - 2,
+        num_batches=len(dummy_binary_db_dataset) // 16 - 2,
         shuffle=True,
         max_num_elements=max_num_elements,
         seed=0,
     )
     refs_seed1 = fit_linear_references(
         ["energy"],
-        dataset=dummy_binary_dataset,
+        dataset=dummy_binary_db_dataset,
         batch_size=16,
-        num_batches=len(dummy_binary_dataset) // 16 - 2,
+        num_batches=len(dummy_binary_db_dataset) // 16 - 2,
         shuffle=True,
         max_num_elements=max_num_elements,
         seed=0,
     )
     refs_noseed = fit_linear_references(
         ["energy"],
-        dataset=dummy_binary_dataset,
+        dataset=dummy_binary_db_dataset,
         batch_size=16,
-        num_batches=len(dummy_binary_dataset) // 16 - 2,
+        num_batches=len(dummy_binary_db_dataset) // 16 - 2,
         shuffle=True,
         max_num_elements=max_num_elements,
         seed=1,

tests/core/modules/test_normalizer.py

@@ -20,18 +20,18 @@
 
 
 @pytest.fixture(scope="session")
-def normalizers(dummy_binary_dataset):
+def normalizers(dummy_binary_db_dataset):
     return fit_normalizers(
         ["energy", "forces"],
         override_values={"forces": {"mean": 0.0}},
-        dataset=dummy_binary_dataset,
+        dataset=dummy_binary_db_dataset,
         batch_size=16,
         shuffle=False,
     )
 
 
-def test_norm_denorm(normalizers, dummy_binary_dataset, dummy_element_refs):
-    batch = data_list_collater(list(dummy_binary_dataset), otf_graph=True)
+def test_norm_denorm(normalizers, dummy_binary_db_dataset, dummy_element_refs):
+    batch = data_list_collater(list(dummy_binary_db_dataset), otf_graph=True)
     # test norm and denorm
     for target, normalizer in normalizers.items():
         normed = normalizer.norm(batch[target])
@@ -43,7 +43,7 @@ def test_norm_denorm(normalizers, dummy_binary_dataset, dummy_element_refs):
         )
 
 
-def test_create_normalizers(normalizers, dummy_binary_dataset, tmp_path):
+def test_create_normalizers(normalizers, dummy_binary_db_dataset, tmp_path):
     # test that forces mean was overriden
     assert normalizers["forces"].mean.item() == 0.0
 
@@ -81,7 +81,7 @@ def test_create_normalizers(normalizers, dummy_binary_dataset, tmp_path):
     assert norm.state_dict() == sdict
 
     # from tensor directly
-    batch = data_list_collater(list(dummy_binary_dataset), otf_graph=True)
+    batch = data_list_collater(list(dummy_binary_db_dataset), otf_graph=True)
     norm = create_normalizer(tensor=batch.energy)
     assert isinstance(norm, Normalizer)
     # assert norm.state_dict() == sdict