make data transforms nn.Modules for device assignment

Author: cw-tan

nequip/data/transforms/neighborlist.py

@@ -5,7 +5,7 @@
 from typing import Optional, Dict, Union, List
 
 
-class NeighborListTransform:
+class NeighborListTransform(torch.nn.Module):
     """Constructs a neighborlist and adds it to the ``AtomicDataDict``.
 
     Args:
@@ -23,6 +23,8 @@ def __init__(
         type_names: Optional[List[str]] = None,
         **kwargs,
     ):
+        super().__init__()
+
         self.r_max = r_max
         self.type_names = type_names
         self.per_edge_type_cutoff = per_edge_type_cutoff
@@ -40,7 +42,7 @@ def __init__(
                 type_names=type_names,
             )
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         data = compute_neighborlist_(data, self.r_max, **self.kwargs)
 
         # prune based on per-edge-type cutoffs if specified
@@ -50,7 +52,7 @@ def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         return data
 
 
-class NeighborListPruneTransform:
+class NeighborListPruneTransform(torch.nn.Module):
     """Prunes a neighborlist based on per-edge-type cutoffs.
 
     Args:
@@ -65,6 +67,8 @@ def __init__(
         per_edge_type_cutoff: Dict[str, Union[float, Dict[str, float]]],
         type_names: List[str],
     ):
+        super().__init__()
+
         self.r_max = r_max
         self.per_edge_type_cutoff = per_edge_type_cutoff
         self.type_names = type_names
@@ -78,7 +82,7 @@ def __init__(
             per_edge_type_cutoff=per_edge_type_cutoff,
         )
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         """Prune neighbor list based on per-edge-type cutoffs."""
 
         if AtomicDataDict.ATOM_TYPE_KEY not in data:
@@ -114,9 +118,9 @@ def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
 class SortedNeighborListTransform(NeighborListTransform):
     """Behaves like :class:`NeighborListTransform` but additionally sorts the neighborlist and provides transpose permutation indices."""
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # first compute the basic neighborlist
-        data = super().__call__(data)
+        data = super().forward(data)
 
         # sort the edge index and corresponding edge attributes
         edge_idxs = data[AtomicDataDict.EDGE_INDEX_KEY]

nequip/data/transforms/stress_utils.py

@@ -3,7 +3,7 @@
 from nequip.data import AtomicDataDict
 
 
-class VirialToStressTransform:
+class VirialToStressTransform(torch.nn.Module):
     r"""Converts virials to stress and adds the stress to the ``AtomicDataDict``.
 
     Specifically implements
@@ -15,9 +15,9 @@ class VirialToStressTransform:
     """
 
     def __init__(self):
-        pass
+        super().__init__()
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # see discussion in https://github.com/libAtoms/QUIP/issues/227 about sign convention
         # they say the standard convention is virial = -stress x volume
         # we assume that the AtomicDataDict contains virials
@@ -29,23 +29,23 @@ def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         return data
 
 
-class StressSignFlipTransform:
+class StressSignFlipTransform(torch.nn.Module):
     r"""Flips the sign of stress in the ``AtomicDataDict``.
 
     In the NequIP convention, positive diagonal components of the stress tensor implies that the system is under tensile strain and wants to compress, while a negative value implies that the system is under compressive strain and wants to expand.
     This transform can be applied to datasets that follow the opposite sign convention, so that the necessary sign flip happens on-the-fly during training and users can avoid having to generate a copy of the dataset with NequIP stress sign conventions.
     """
 
     def __init__(self):
-        pass
+        super().__init__()
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # see discussion in https://github.com/libAtoms/QUIP/issues/227 about sign convention
         data[AtomicDataDict.STRESS_KEY] = data[AtomicDataDict.STRESS_KEY].neg()
         return data
 
 
-class AddNaNStressTransform:
+class AddNaNStressTransform(torch.nn.Module):
     """Add NaN stress tensors for structures without stress data.
 
     Useful for datasets where stresses are not available for all structures.
@@ -54,9 +54,9 @@ class AddNaNStressTransform:
     """
 
     def __init__(self):
-        pass
+        super().__init__()
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # only add if stress is not already present
         if AtomicDataDict.STRESS_KEY not in data:
             num_frames = AtomicDataDict.num_frames(data)

nequip/data/transforms/type_mapper.py

@@ -7,7 +7,7 @@
 from typing import List, Dict, Optional
 
 
-class ChemicalSpeciesToAtomTypeMapper:
+class ChemicalSpeciesToAtomTypeMapper(torch.nn.Module):
     """Maps atomic numbers to atom types and adds the atom types to the ``AtomicDataDict``.
 
     This transform accounts for how the atom types seen by the model can be different from the atomic species that one obtains from a conventional dataset. There could be cases where the same chemical species corresponds to multiple atom types, e.g. different charge states.
@@ -23,6 +23,8 @@ def __init__(
         chemical_species_to_atom_type_map: Optional[Dict[str, str]] = None,
         chemical_symbols: Optional[List[str]] = None,
     ):
+        super().__init__()
+
         # TODO: eventually remove all this logic
         # error out with deprecated API usage
         if chemical_symbols is not None:
@@ -68,7 +70,7 @@ def __init__(
                 )
 
         # make a lookup table mapping atomic numbers to 0-based model type indexes
-        self.lookup_table = torch.full(
+        lookup_table = torch.full(
             (max(ase.data.atomic_numbers.values()),), -1, dtype=torch.long
         )
 
@@ -78,9 +80,11 @@ def __init__(
                 raise ValueError(f"Invalid chemical symbol '{chem_symbol}'")
             atomic_num = ase.data.atomic_numbers[chem_symbol]
             type_idx = type_name_to_index[atom_type_name]
-            self.lookup_table[atomic_num] = type_idx
+            lookup_table[atomic_num] = type_idx
+
+        self.register_buffer("lookup_table", lookup_table)
 
-    def __call__(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
+    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         if AtomicDataDict.ATOM_TYPE_KEY in data:
             raise RuntimeError(f"Data already contains {AtomicDataDict.ATOM_TYPE_KEY}")