diff --git a/src/constants.py b/src/constants.py index f0f47a5..1c0d2d2 100644 --- a/src/constants.py +++ b/src/constants.py @@ -4,4 +4,4 @@ REPO_PATH = pathlib.Path(__file__).parents[1] -DBAND_FILE_PATH = REPO_PATH / "data" / "dband_centers.csv" \ No newline at end of file +DBAND_FILE_PATH = REPO_PATH / "data" / "dband_centers.csv" diff --git a/src/featurizers.py b/src/featurizers.py index 04b9ac7..b55d7a1 100644 --- a/src/featurizers.py +++ b/src/featurizers.py @@ -3,18 +3,19 @@ import abc import csv -import numpy as np import networkx as nx +import numpy as np import torch - -from torch.nn.functional import one_hot from mendeleev import element +from torch.nn.functional import one_hot from constants import DBAND_FILE_PATH from graphs import AtomsGraph + class OneHotEncoder: """Featurize a property using a one-hot encoding scheme.""" + def __init__(self): """Blank constructor.""" pass @@ -37,7 +38,7 @@ def fit(self, min, max, n_intervals): def transform(self, property): """Transform a given property vector/matrix/tensor. - + Parameters ---------- property: list or np.ndarray or torch.Tensor @@ -48,24 +49,24 @@ def transform(self, property): property = torch.Tensor(property) # Scale between 0 and num_intervals - scaled_prop = ((property - self.min) - / (self.max - self.min)) * self.n_intervals - + scaled_prop = ((property - self.min) / (self.max - self.min)) * self.n_intervals + # Apply floor function floor_prop = torch.floor(scaled_prop) - + # Create onehot encoding - onehot_prop = one_hot(floor_prop.to(torch.int64), - num_classes=self.n_intervals) - + onehot_prop = one_hot(floor_prop.to(torch.int64), num_classes=self.n_intervals) + return onehot_prop + class Featurizer(abc.ABC): """Meta class for defining featurizers.""" + @abc.abstractmethod def __init__(self, encoder): """Initialize class variables and fit encoder. - + Parameters ---------- encoder: OneHotEncoder @@ -92,7 +93,7 @@ def featurize_graph(self, atoms_graph): @abc.abstractproperty def feat_tensor(self): """Return the featurized node tensor. - + Returns ------- feat_tensor: torch.Tensor @@ -100,11 +101,11 @@ def feat_tensor(self): M = n_intervals provided to the encoder """ pass - + @abc.abstractstaticmethod def name(self): """Return the name of the featurizer. - + Returns ------- _name = str @@ -115,6 +116,7 @@ def name(self): class AtomNumFeaturizer(Featurizer): """Featurize nodes based on atomic number.""" + def __init__(self, encoder, min=0, max=80, n_intervals=10): """Initialize featurizer with min = 0, max = 80, n_intervals = 10. @@ -156,7 +158,7 @@ def featurize_graph(self, atoms_graph): @property def feat_tensor(self): """Return the featurized node tensor. - + Returns ------- feat_tensor: torch.Tensor @@ -164,14 +166,16 @@ def feat_tensor(self): M = n_intervals provided to the encoder """ return self._feat_tensor - + @staticmethod def name(): """Return the name of the featurizer.""" return "atomic_number" + class DBandFeaturizer(Featurizer): """Featurize nodes based on close-packed d-band center.""" + def __init__(self, encoder, min=-5, max=3, n_intervals=10): """Initialize featurizer with min = -5, max = 3, n_intervals = 10. @@ -223,7 +227,7 @@ def featurize_graph(self, atoms_graph): @property def feat_tensor(self): """Return the featurized node tensor. - + Returns ------- feat_tensor: torch.Tensor @@ -236,9 +240,11 @@ def feat_tensor(self): def name(): """Return the name of the featurizer.""" return "dband_center" - + + class ValenceFeaturizer(Featurizer): """Featurize nodes based on number of valence electrons.""" + def __init__(self, encoder, min=1, max=12, n_intervals=12): """Initialize featurizer with min = 1, max = 12, n_intervals = 12. @@ -263,7 +269,7 @@ def __init__(self, encoder, min=1, max=12, n_intervals=12): self.encoder.fit(self.min, self.max, self.n_intervals) # Create a map between atomic number and number of valence electrons - self.map_dict = {0: 0, 1: 1, 2:0} + self.map_dict = {0: 0, 1: 1, 2: 0} for i in range(3, 21, 1): self.map_dict[i] = min(element(i).ec.get_valence().ne(), 12) @@ -285,7 +291,7 @@ def featurize_graph(self, atoms_graph): @property def feat_tensor(self): """Return the featurized node tensor. - + Returns ------- feat_tensor: torch.Tensor @@ -298,9 +304,11 @@ def feat_tensor(self): def name(): """Return the name of the featurizer.""" return "valence" - + + class CoordinationFeaturizer(Featurizer): """Featurize nodes based on coordination number.""" + def __init__(self, encoder, min=1, max=15, n_intervals=15): """Initialize featurizer with min = 1, max = 15, n_intervals = 15. @@ -342,7 +350,7 @@ def featurize_graph(self, atoms_graph): @property def feat_tensor(self): """Return the featurized node tensor. - + Returns ------- feat_tensor: torch.Tensor @@ -355,9 +363,11 @@ def feat_tensor(self): def name(): """Return the name of the featurizer.""" return "coordination" - + + class BondDistanceFeaturizer(Featurizer): """Featurize edges based on bond distance.""" + def __init__(self, encoder, min, max, n_intervals): """Initialize bond distance featurizer. @@ -395,14 +405,14 @@ def featurize_graph(self, atoms_graph): # Create node feature matrix self._feat_tensor = self.encoder.transform(bond_dist_arr) - + # Create list of edge indices self._edge_indices = torch.Tensor(list(atoms_graph.graph.edges())) @property def feat_tensor(self): """Return the featurized node tensor. - + Returns ------- feat_tensor: torch.Tensor @@ -410,7 +420,7 @@ def feat_tensor(self): M = n_intervals provided to the encoder """ return self._feat_tensor - + @property def edge_indices(self): """Return list of edge indices. @@ -426,48 +436,55 @@ def edge_indices(self): def name(): """Return the name of the featurizer.""" return "valence" - + + class BulkBondDistanceFeaturizer(BondDistanceFeaturizer): """Featurize bulk bond distances. - + Child class of BondDistanceFeaturizer with suitable min, max, and n_interval values initialized for bulk atoms. The values are: min = 0, max = 8, n_intervals = 8. """ + def __init__(self, encoder, min=0, max=8, n_intervals=8): super().__init__(encoder, min=min, max=max, n_intervals=n_intervals) + class SurfaceBondDistanceFeaturizer(BondDistanceFeaturizer): """Featurize bulk bond distances. - + Child class of BondDistanceFeaturizer with suitable min, max, and n_interval values initialized for surface atoms. The values are: min = 0, max = 5, n_intervals = 10. """ + def __init__(self, encoder, min=0, max=5, n_intervals=10): super().__init__(encoder, min=min, max=max, n_intervals=n_intervals) + class AdsorbateBondDistanceFeaturizer(BondDistanceFeaturizer): """Featurize bulk bond distances. - + Child class of BondDistanceFeaturizer with suitable min, max, and n_interval values initialized for adsorbate atoms. The values are: min = 0, max = 4, n_intervals = 16. """ + def __init__(self, encoder, min=0, max=4, n_intervals=16): super().__init__(encoder, min=min, max=max, n_intervals=n_intervals) + if __name__ == "__main__": from ase.io import read atoms = read("CONTCAR") g = AtomsGraph(atoms, select_idx=[1, 10, 11, 12]) - + anf = AtomNumFeaturizer(OneHotEncoder()) anf.featurize_graph(g) print(anf.feat_tensor.shape) - + bdf = BulkBondDistanceFeaturizer(OneHotEncoder()) bdf.featurize_graph(g) print(bdf.feat_tensor) - print(bdf.edge_indices) \ No newline at end of file + print(bdf.edge_indices) diff --git a/src/graphs.py b/src/graphs.py index 99e96bc..b744e8f 100644 --- a/src/graphs.py +++ b/src/graphs.py @@ -4,8 +4,8 @@ import networkx as nx import numpy as np -from ase.neighborlist import (build_neighbor_list, natural_cutoffs, - NewPrimitiveNeighborList) +from ase.neighborlist import (NewPrimitiveNeighborList, build_neighbor_list, + natural_cutoffs) class AtomsGraph: @@ -43,8 +43,9 @@ def create_graph(self): self.node_count = 0 self.map_idx_node = {} for atom in self.atoms: - if (atom.index in self.select_idx - and atom.index not in list(graph.nodes(data="index"))): + if atom.index in self.select_idx and atom.index not in list( + graph.nodes(data="index") + ): graph.add_node( self.node_count, index=atom.index, @@ -60,7 +61,7 @@ def create_graph(self): natural_cutoffs(self.atoms), bothways=True, self_interaction=False, - primitive=NewPrimitiveNeighborList + primitive=NewPrimitiveNeighborList, ) # Iterate over nodes, identify neighbors, and add edges between them @@ -101,9 +102,7 @@ def create_graph(self): self.atoms[nn].position, offset, ) - graph.add_edge( - n, self.map_idx_node[nn], bond_distance=bond_dist - ) + graph.add_edge(n, self.map_idx_node[nn], bond_distance=bond_dist) # Pad graph graph = self.pad_graph(graph) @@ -120,12 +119,12 @@ def pad_graph(self, graph): This can be used to make sure that the number of nodes in each graph is equal to max_atoms - + Parameters ---------- graph: Networkx.Graph A Networkx graph - + Returns ------- padded_graph: Networkx.Graph @@ -135,11 +134,7 @@ def pad_graph(self, graph): for i in range(self.node_count, self.max_atoms, 1): padded_graph.add_node( - i, - index=-1, - atomic_number=0, - symbol="", - position=np.zeros(3) + i, index=-1, atomic_number=0, symbol="", position=np.zeros(3) ) return padded_graph @@ -159,13 +154,13 @@ def calc_minimum_distance(self, pos_1, pos_2, offset): """ # First, calculate the distance without offset dist_1 = np.linalg.norm(pos_1 - pos_2) - + # Next calculate the distance by applying offset to second position dist_2 = np.linalg.norm(pos_1 - (pos_2 + offset @ self.atoms.get_cell())) # Get minimum distance min_dist = min(dist_1, dist_2) - + return min_dist def plot(self, filename=None): @@ -178,10 +173,11 @@ def plot(self, filename=None): """ pass + if __name__ == "__main__": from ase.io import read - + atoms = read("CONTCAR") g = AtomsGraph(atoms, select_idx=[1, 10, 11, 12]) print(g.map_idx_node) - print(g.graph.edges(data="bond_distance")) \ No newline at end of file + print(g.graph.edges(data="bond_distance"))