Setting the default directory for storing MGLDataset in the current working directory.

examples/Training a M3GNet Potential with PyTorch Lightning.ipynb

@@ -127,7 +127,11 @@
     "element_types = DEFAULT_ELEMENTS\n",
     "converter = Structure2Graph(element_types=element_types, cutoff=5.0)\n",
     "dataset = MGLDataset(\n",
-    "    threebody_cutoff=4.0, structures=structures, converter=converter, labels=labels, include_line_graph=True\n",
+    "    threebody_cutoff=4.0, \n",
+    "    structures=structures, \n",
+    "    converter=converter, \n",
+    "    labels=labels, \n",
+    "    include_line_graph=True,\n",
     ")\n",
     "train_data, val_data, test_data = split_dataset(\n",
     "    dataset,\n",

src/matgl/ext/ase.py

@@ -26,6 +26,7 @@
 from ase.md.nvtberendsen import NVTBerendsen
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.md.verlet import VelocityVerlet
+from ase.stress import full_3x3_to_voigt_6_stress
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 from pymatgen.optimization.neighbors import find_points_in_spheres
@@ -133,6 +134,7 @@ def __init__(
         potential: Potential,
         state_attr: torch.Tensor | None = None,
         stress_weight: float = 1.0,
+        use_voigt: bool = False,
         **kwargs,
     ):
         """
@@ -143,6 +145,7 @@ def __init__(
             state_attr (tensor): State attribute
             compute_stress (bool): whether to calculate the stress
             stress_weight (float): conversion factor from GPa to eV/A^3, if it is set to 1.0, the unit is in GPa
+            use_voigt (bool): whether the voigt notation is used for stress output
             **kwargs: Kwargs pass through to super().__init__().
         """
         super().__init__(**kwargs)
@@ -154,6 +157,7 @@ def __init__(
         self.state_attr = state_attr
         self.element_types = potential.model.element_types  # type: ignore
         self.cutoff = potential.model.cutoff
+        self.use_voigt = use_voigt
 
     def calculate(
         self,
@@ -186,7 +190,12 @@ def calculate(
             forces=calc_result[1].detach().cpu().numpy(),
         )
         if self.compute_stress:
-            self.results.update(stress=calc_result[2].detach().cpu().numpy() * self.stress_weight)
+            stresses_np = (
+                full_3x3_to_voigt_6_stress(calc_result[2].detach().cpu().numpy())
+                if self.use_voigt
+                else calc_result[2].detach().cpu().numpy()
+            )
+            self.results.update(stress=stresses_np * self.stress_weight)
         if self.compute_hessian:
             self.results.update(hessian=calc_result[3].detach().cpu().numpy())
         if self.compute_magmom:

src/matgl/graph/data.py

@@ -130,12 +130,12 @@ def __init__(
         directed_line_graph: bool = False,
         structures: list | None = None,
         labels: dict[str, list] | None = None,
-        name: str = "MGLDataset",
+        directory_name: str = "MGLDataset",
         graph_labels: list[int | float] | None = None,
         clear_processed: bool = False,
         save_cache: bool = True,
-        raw_dir: str | None = None,
-        save_dir: str | None = None,
+        raw_dir: str = "./",
+        save_dir: str = "./",
     ):
         """
         Args:
@@ -152,7 +152,7 @@ def __init__(
                 Default: False (for M3GNet)
             structures: Pymatgen structure.
             labels: targets, as a dict of {name: list of values}.
-            name: name of dataset.
+            directory_name: name of the generated directory that stores the dataset.
             graph_labels: state attributes.
             clear_processed: Whether to clear the stored structures after processing into graphs. Structures
                 are not really needed after the conversion to DGL graphs and can take a significant amount of memory.
@@ -161,7 +161,7 @@ def __init__(
                 Default: True
             raw_dir : str specifying the directory that will store the downloaded data or the directory that already
                 stores the input data.
-                Default: ~/.dgl/
+                Default: current working directory
             save_dir : directory to save the processed dataset. Default: same as raw_dir.
         """
         self.filename = filename
@@ -180,7 +180,7 @@ def __init__(
         self.graph_labels = graph_labels
         self.clear_processed = clear_processed
         self.save_cache = save_cache
-        super().__init__(name=name, raw_dir=raw_dir, save_dir=save_dir)
+        super().__init__(name=directory_name, raw_dir=raw_dir, save_dir=save_dir)
 
     def has_cache(self) -> bool:
         """Check if the dgl_graph.bin exists or not."""

tests/graph/test_data.py

@@ -72,7 +72,7 @@ def test_megnet_dataset_with_graph_label_float(self, LiFePO4, BaNiO3):
             labels={"label": label},
             clear_processed=True,
             graph_labels=graph_label,
-            name="MGLDataset_megnet",
+            directory_name="MGLDataset_megnet",
         )
         g1, lat1, state1, label1 = dataset_with_graph_label[0]
         g2, lat2, state2, label2 = dataset_with_graph_label[1]
@@ -85,7 +85,7 @@ def test_load_megenet_dataset(self, LiFePO4, BaNiO3):
         label = [-1.0, 2.0]
         element_types = get_element_list(structures)
         cry_graph = Structure2Graph(element_types=element_types, cutoff=4.0)
-        dataset = MGLDataset(name="MGLDataset_megnet")
+        dataset = MGLDataset(directory_name="MGLDataset_megnet")
         g1, lat1, state1, label1 = dataset[0]
         assert label1["label"] == label[0]
         assert g1.num_edges() == cry_graph.get_graph(LiFePO4)[0].num_edges()
@@ -127,7 +127,7 @@ def test_mgl_dataset(self, LiFePO4, BaNiO3):
             include_line_graph=True,
             labels={"energies": energies, "forces": forces, "stresses": stresses},
             clear_processed=True,
-            name="MGLDataset_pes",
+            directory_name="MGLDataset_pes",
         )
         g1, lat1, l_g1, state1, pes1 = dataset[0]
         g2, lat2, l_g2, state2, pes2 = dataset[1]
@@ -148,7 +148,7 @@ def test_load_mgl_dataset(self, LiFePO4, BaNiO3):
         element_types = get_element_list(structures)
         cry_graph = Structure2Graph(element_types=element_types, cutoff=4.0)
         dataset = MGLDataset(
-            name="MGLDataset_pes",
+            directory_name="MGLDataset_pes",
             include_line_graph=True,
         )
         dataset.load()