Merge pull request #9 from GreeleyGroup/enh/utils

ENH: Added util functions

data/dband_centers.csv

@@ -1,3 +1,4 @@
+0, 0
 22, 1.50
 23, 1.06
 24, 0.16

src/featurizers.py

@@ -151,9 +151,11 @@ def featurize_graph(self, atoms_graph):
         # Get atomic numbers
         atom_num_dict = nx.get_node_attributes(atoms_graph.graph, "atomic_number")
         atom_num_arr = np.array(list(atom_num_dict.values()))
+        zero_idx = np.argwhere(atom_num_arr == 0.0)
 
         # Create node feature matrix
         self._feat_tensor = self.encoder.transform(atom_num_arr)
+        self._feat_tensor[zero_idx, :] = 0.0
 
     @property
     def feat_tensor(self):
@@ -217,12 +219,14 @@ def featurize_graph(self, atoms_graph):
         # Get atomic numbers
         atom_num_dict = nx.get_node_attributes(atoms_graph.graph, "atomic_number")
         atom_num_arr = np.array(list(atom_num_dict.values()))
+        zero_idx = np.argwhere(atom_num_arr == 0.0)
 
         # Map from atomic number to d-band center
         dband_arr = np.vectorize(self.map_dict.__getitem__)(atom_num_arr)
 
         # Create node feature matrix
         self._feat_tensor = self.encoder.transform(dband_arr)
+        self._feat_tensor[zero_idx, :] = 0.0
 
     @property
     def feat_tensor(self):
@@ -284,9 +288,14 @@ def featurize_graph(self, atoms_graph):
         # Get atomic numbers
         atom_num_dict = nx.get_node_attributes(atoms_graph.graph, "atomic_number")
         atom_num_arr = np.array(list(atom_num_dict.values()))
+        zero_idx = np.argwhere(atom_num_arr == 0.0)
+
+        # Get valence electrons for each atom
+        valence_arr = np.vectorize(self.map_dict.__getitem__)(atom_num_arr)
 
         # Create node feature matrix
-        self._feat_tensor = self.encoder.transform(atom_num_arr)
+        self._feat_tensor = self.encoder.transform(valence_arr)
+        self._feat_tensor[zero_idx, :] = 0.0
 
     @property
     def feat_tensor(self):
@@ -306,6 +315,91 @@ def name():
         return "valence"
 
 
+class ReactivityFeaturizer(Featurizer):
+    """Featurize nodes based on close-packed d-band center and/or valence."""
+
+    def __init__(self, encoder, min=-5, max=3, n_intervals=10):
+        """Initialize featurizer with min = -5, max = 3, n_intervals = 10.
+
+        Parameters
+        ----------
+        encoder: OneHotEncoder
+            Initialized object of class OneHotEncoder.
+        min: int
+            Minimum value of d-band center
+        max: int
+            Maximum value of d-band center
+        n_intervals: int
+            Number of intervals
+        """
+        # Initialize variables
+        self.min = min
+        self.max = max
+        self.n_intervals = n_intervals
+
+        # Fit encoder
+        self.encoder_dband = encoder
+        self.encoder_dband.fit(self.min, self.max, self.n_intervals)
+
+        # Fit valence encoder
+        self.encoder_val = encoder
+        self.encoder_val.fit(1, 12, self.n_intervals)
+
+        # Get dband centers from csv
+        self.map_dict_dband = {}
+        with open(DBAND_FILE_PATH, "r") as f:
+            csv_reader = csv.reader(f)
+            for row in csv_reader:
+                self.map_dict_dband[int(row[0])] = float(row[1])
+
+        # Create a map between atomic number and number of valence electrons
+        self.map_dict_val = {0: 0, 1: 1, 2: 0}
+        for i in range(3, 21, 1):
+            self.map_dict_val[i] = min(element(i).ec.get_valence().ne(), 12)
+
+    def featurize_graph(self, atoms_graph):
+        """Featurize an AtomsGraph.
+
+        Parameters
+        ----------
+        graph: AtomsGraph
+            A graph of a collection of bulk, surface, or adsorbate atoms.
+        """
+        # Get atomic numbers
+        atom_num_dict = nx.get_node_attributes(atoms_graph.graph, "atomic_number")
+        atom_num_arr = np.array(list(atom_num_dict.values()))
+        zero_idx = np.argwhere(atom_num_arr == 0.0)
+
+        # Map from atomic number to d-band center
+        react_arr = np.zeros_like(atom_num_arr)
+        for i, n in enumerate(atom_num_arr):
+            if n < 21:
+                react_arr[i] = self.map_dict_val[n]
+            else:
+                react_arr[i] = self.map_dict_dband[n]
+
+        # Create node feature matrix
+        self._feat_tensor = self.encoder.transform(react_arr)
+        self._feat_tensor[zero_idx, :] = 0.0
+
+    @property
+    def feat_tensor(self):
+        """Return the featurized node tensor.
+
+        Returns
+        -------
+        feat_tensor: torch.Tensor
+            Featurized tensor having shape (N, M) where N = number of atoms and
+            M = n_intervals provided to the encoder
+        """
+        return self._feat_tensor
+
+    @staticmethod
+    def name():
+        """Return the name of the featurizer."""
+        return "reactivity"
+
+
 class CoordinationFeaturizer(Featurizer):
     """Featurize nodes based on coordination number."""
 
@@ -341,11 +435,17 @@ def featurize_graph(self, atoms_graph):
             A graph of a collection of bulk, surface, or adsorbate atoms.
         """
         # Get atomic numbers
+        atom_num_dict = nx.get_node_attributes(atoms_graph.graph, "atomic_number")
+        atom_num_arr = np.array(list(atom_num_dict.values()))
+        zero_idx = np.argwhere(atom_num_arr == 0.0)
+
+        # Get coordination numbers
         cn_dict = nx.get_node_attributes(atoms_graph.graph, "coordination")
         cn_arr = np.array(list(cn_dict.values()))
 
         # Create node feature matrix
         self._feat_tensor = self.encoder.transform(cn_arr)
+        self._feat_tensor[zero_idx, :] = 0.0
 
     @property
     def feat_tensor(self):
@@ -407,7 +507,9 @@ def featurize_graph(self, atoms_graph):
         self._feat_tensor = self.encoder.transform(bond_dist_arr)
 
         # Create list of edge indices
-        self._edge_indices = torch.Tensor(list(atoms_graph.graph.edges()))
+        self._edge_indices = torch.LongTensor(list(atoms_graph.graph.edges())).view(
+            2, -1
+        )
 
     @property
     def feat_tensor(self):
@@ -435,7 +537,7 @@ def edge_indices(self):
     @staticmethod
     def name():
         """Return the name of the featurizer."""
-        return "valence"
+        return "bond_distance"
 
 
 class BulkBondDistanceFeaturizer(BondDistanceFeaturizer):
@@ -449,6 +551,11 @@ class BulkBondDistanceFeaturizer(BondDistanceFeaturizer):
     def __init__(self, encoder, min=0, max=8, n_intervals=8):
         super().__init__(encoder, min=min, max=max, n_intervals=n_intervals)
 
+    @staticmethod
+    def name():
+        """Return the name of the featurizer."""
+        return "bulk_bond_distance"
+
 
 class SurfaceBondDistanceFeaturizer(BondDistanceFeaturizer):
     """Featurize bulk bond distances.
@@ -461,6 +568,11 @@ class SurfaceBondDistanceFeaturizer(BondDistanceFeaturizer):
     def __init__(self, encoder, min=0, max=5, n_intervals=10):
         super().__init__(encoder, min=min, max=max, n_intervals=n_intervals)
 
+    @staticmethod
+    def name():
+        """Return the name of the featurizer."""
+        return "surface_bond_distance"
+
 
 class AdsorbateBondDistanceFeaturizer(BondDistanceFeaturizer):
     """Featurize bulk bond distances.
@@ -473,6 +585,24 @@ class AdsorbateBondDistanceFeaturizer(BondDistanceFeaturizer):
     def __init__(self, encoder, min=0, max=4, n_intervals=16):
         super().__init__(encoder, min=min, max=max, n_intervals=n_intervals)
 
+    @staticmethod
+    def name():
+        """Return the name of the featurizer."""
+        return "adsorbate_distance"
+
+
+list_of_node_featurizers = [
+    AtomNumFeaturizer,
+    DBandFeaturizer,
+    ValenceFeaturizer,
+    CoordinationFeaturizer,
+]
+
+list_of_edge_featurizers = [
+    BulkBondDistanceFeaturizer,
+    SurfaceBondDistanceFeaturizer,
+    AdsorbateBondDistanceFeaturizer,
+]
 
 if __name__ == "__main__":
     from ase.io import read
@@ -486,5 +616,5 @@ def __init__(self, encoder, min=0, max=4, n_intervals=16):
 
     bdf = BulkBondDistanceFeaturizer(OneHotEncoder())
     bdf.featurize_graph(g)
-    print(bdf.feat_tensor)
+    print(bdf.feat_tensor.shape)
     print(bdf.edge_indices)

src/graphs.py

@@ -4,8 +4,11 @@
 
 import networkx as nx
 import numpy as np
-from ase.neighborlist import (NewPrimitiveNeighborList, build_neighbor_list,
-                              natural_cutoffs)
+from ase.neighborlist import (
+    NewPrimitiveNeighborList,
+    build_neighbor_list,
+    natural_cutoffs,
+)
 
 
 class AtomsGraph: