Fix codestyle

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,15 +437,15 @@ 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")
         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.
+        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):

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:

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)
-