From f1b88eec6f6468e2a230f37d45b88b53a286cb2a Mon Sep 17 00:00:00 2001
From: Janosh Riebesell <janosh.riebesell@gmail.com>
Date: Mon, 11 Mar 2024 16:02:18 +0100
Subject: [PATCH] refactor CrystalGraphConverter._create_graph_fast

---
 chgnet/graph/converter.py | 20 +++++++-------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/chgnet/graph/converter.py b/chgnet/graph/converter.py
index 9a29a856..376dcb5b 100644
--- a/chgnet/graph/converter.py
+++ b/chgnet/graph/converter.py
@@ -250,25 +250,19 @@ def _create_graph_fast(
         distance = np.ascontiguousarray(distance)
         gc_saved = gc.get_threshold()
         gc.set_threshold(0)
-        (
-            nodes,
-            directed_edges_list,
-            undirected_edges_list,
-            undirected_edges,
-        ) = make_graph(
+        nodes, dir_edges_list, undir_edges_list, undirected_edges = make_graph(
             center_index, len(center_index), neighbor_index, image, distance, n_atoms
         )
         graph = Graph(nodes=nodes)
-        graph.directed_edges_list = directed_edges_list
-        graph.undirected_edges_list = undirected_edges_list
+        graph.directed_edges_list = dir_edges_list
+        graph.undirected_edges_list = undir_edges_list
         graph.undirected_edges = undirected_edges
         gc.set_threshold(gc_saved[0])
 
         return graph
 
     def set_isolated_atom_response(
-        self,
-        on_isolated_atoms: Literal["ignore", "warn", "error"],
+        self, on_isolated_atoms: Literal["ignore", "warn", "error"]
     ) -> None:
         """Set the graph converter's response to isolated atom graph
         Args:
@@ -282,7 +276,7 @@ def set_isolated_atom_response(
         self.on_isolated_atoms = on_isolated_atoms
         return
 
-    def as_dict(self) -> dict[str, float]:
+    def as_dict(self) -> dict[str, str | float]:
         """Save the args of the graph converter."""
         return {
             "atom_graph_cutoff": self.atom_graph_cutoff,
@@ -291,6 +285,6 @@ def as_dict(self) -> dict[str, float]:
         }
 
     @classmethod
-    def from_dict(cls, dict) -> CrystalGraphConverter:
+    def from_dict(cls, dct: dict) -> CrystalGraphConverter:
         """Create converter from dictionary."""
-        return CrystalGraphConverter(**dict)
+        return CrystalGraphConverter(**dct)