diff --git a/src/featurizers.py b/src/featurizers.py index 9597295..145152f 100644 --- a/src/featurizers.py +++ b/src/featurizers.py @@ -151,11 +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.) + 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. + self._feat_tensor[zero_idx, :] = 0.0 @property def feat_tensor(self): @@ -219,14 +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.) + 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. + self._feat_tensor[zero_idx, :] = 0.0 @property def feat_tensor(self): @@ -288,14 +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.) + 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(valence_arr) - self._feat_tensor[zero_idx, :] = 0. + self._feat_tensor[zero_idx, :] = 0.0 @property def feat_tensor(self): @@ -314,6 +314,7 @@ def name(): """Return the name of the featurizer.""" return "valence" + class ReactivityFeaturizer(Featurizer): """Featurize nodes based on close-packed d-band center and/or valence.""" @@ -367,7 +368,7 @@ 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.) + 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) @@ -379,7 +380,7 @@ def featurize_graph(self, atoms_graph): # Create node feature matrix self._feat_tensor = self.encoder.transform(react_arr) - self._feat_tensor[zero_idx, :] = 0. + self._feat_tensor[zero_idx, :] = 0.0 @property def feat_tensor(self): @@ -398,6 +399,7 @@ def name(): """Return the name of the featurizer.""" return "reactivity" + class CoordinationFeaturizer(Featurizer): """Featurize nodes based on coordination number.""" @@ -435,7 +437,7 @@ 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.) + zero_idx = np.argwhere(atom_num_arr == 0.0) # Get coordination numbers cn_dict = nx.get_node_attributes(atoms_graph.graph, "coordination") @@ -443,7 +445,7 @@ def featurize_graph(self, atoms_graph): # Create node feature matrix self._feat_tensor = self.encoder.transform(cn_arr) - self._feat_tensor[zero_idx, :] = 0. + self._feat_tensor[zero_idx, :] = 0.0 @property def feat_tensor(self): @@ -505,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.LongTensor(list(atoms_graph.graph.edges())).view(2, -1) + self._edge_indices = torch.LongTensor(list(atoms_graph.graph.edges())).view( + 2, -1 + ) @property def feat_tensor(self): diff --git a/src/graphs.py b/src/graphs.py index b744e8f..40a09b3 100644 --- a/src/graphs.py +++ b/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: diff --git a/src/utils.py b/src/utils.py index 1f2fdce..3ecf32b 100644 --- a/src/utils.py +++ b/src/utils.py @@ -6,15 +6,16 @@ import torch from featurizers import ( - list_of_node_featurizers, + OneHotEncoder, list_of_edge_featurizers, - OneHotEncoder + list_of_node_featurizers, ) from graphs import AtomsGraph + def partition_structure(atoms, n_partitions, z_cutoffs): """Partition atomic structue into bulk, surface, and/or adsorbates. - + Parameters ---------- atoms: ase.Atoms object @@ -24,13 +25,15 @@ def partition_structure(atoms, n_partitions, z_cutoffs): z_cutoffs: list or np.ndarray List of z-coordinate cutoffs. xy planes are placed at the specified cutoffs to partition atoms above and below them. The length of z-cutoffs - should be equal to one less than the number of partitions. + should be equal to one less than the number of partitions. """ # Check if length of z_cutoffs is equal to n_paritions if len(z_cutoffs) != n_partitions - 1: - raise ValueError("The length of z_cutoffs must be equal to\ - one less than the number of partitions") - + raise ValueError( + "The length of z_cutoffs must be equal to\ + one less than the number of partitions" + ) + # Add 0 and infinity to cutoffs z_cutoffs = np.insert(z_cutoffs, 0, 0) z_cutoffs = np.insert(z_cutoffs, len(z_cutoffs), np.inf) @@ -41,21 +44,24 @@ def partition_structure(atoms, n_partitions, z_cutoffs): # Iterate over number of partitions part_atoms = [] for i in range(n_partitions): - part_idx = np.argwhere( - (pos[:, -1] >= z_cutoffs[i]) & (pos[:, -1] < z_cutoffs[i+1]) - ).flatten().tolist() + part_idx = ( + np.argwhere((pos[:, -1] >= z_cutoffs[i]) & (pos[:, -1] < z_cutoffs[i + 1])) + .flatten() + .tolist() + ) part_atoms.append(part_idx) return part_atoms + def featurize_atoms( - atoms, - select_idx, - node_features, - edge_features, - max_atoms=50, - encoder=OneHotEncoder() - ): + atoms, + select_idx, + node_features, + edge_features, + max_atoms=50, + encoder=OneHotEncoder(), +): """Featurize atoms and bonds with the chosen featurizers. Parameters @@ -128,18 +134,26 @@ def featurize_atoms( # Store edge indices edge_indices = ef.edge_indices - return {"node_tensor": node_tensor, "edge_tensor": edge_tensor, - "edge_indices": edge_indices} + return { + "node_tensor": node_tensor, + "edge_tensor": edge_tensor, + "edge_indices": edge_indices, + } if __name__ == "__main__": from ase.io import read + atoms = read("CONTCAR") part_atoms = partition_structure(atoms, 3, z_cutoffs=[15, 23.5]) print(part_atoms) - feat_dict = featurize_atoms(atoms, part_atoms[0], ["atomic_number", "dband_center"], - ["bulk_bond_distance"], max_atoms=34) + feat_dict = featurize_atoms( + atoms, + part_atoms[0], + ["atomic_number", "dband_center"], + ["bulk_bond_distance"], + max_atoms=34, + ) print(feat_dict) - \ No newline at end of file