Merge pull request #142 from melo-gonzo/nomad-devset-dataset-update

Nomad Dataset and Devset Update

examples/datasets/nomad/single_task_devset.py

@@ -3,7 +3,10 @@
 import pytorch_lightning as pl
 from torch.nn import LayerNorm, SiLU
 
-from matsciml.datasets.transforms import PointCloudToGraphTransform
+from matsciml.datasets.transforms import (
+    PointCloudToGraphTransform,
+    PeriodicPropertiesTransform,
+)
 from matsciml.lightning.data_utils import MatSciMLDataModule
 from matsciml.models import PLEGNNBackbone
 from matsciml.models.base import ScalarRegressionTask
@@ -49,14 +52,15 @@
         "input_dim": 128,
     },
     lr=1e-3,
-    task_keys=["energy_total"],
+    task_keys=["relative_energy"],
 )
 
 # configure materials project from devset
 dm = MatSciMLDataModule.from_devset(
     "NomadDataset",
     dset_kwargs={
         "transforms": [
+            PeriodicPropertiesTransform(cutoff_radius=6.0, adaptive_cutoff=True),
             PointCloudToGraphTransform(
                 "dgl",
                 cutoff_dist=20.0,

matsciml/datasets/nomad/dataset.py

@@ -44,7 +44,7 @@ def target_keys(self) -> dict[str, list[str]]:
             Dict[str, List[str]]: target keys
         """
         return {
-            "regression": ["energy_total", "efermi"],
+            "regression": ["relative_energy", "energy_total", "efermi"],
             "classification": ["spin_polarized"],
         }
 
@@ -205,5 +205,7 @@ def data_from_key(self, lmdb_index: int, subindex: int) -> Any:
         # data['energies'].keys()
         data = super().data_from_key(lmdb_index, subindex)
         return_dict = {}
+        for k in ["reference_energy", "reference_structure", "relative_energy"]:
+            return_dict[k] = data[k]
         self._parse_data(data, return_dict=return_dict)
         return return_dict