feat: add get_structure_graph helper

kjappelbaum · Jul 28, 2022 · 53fceb4 · 53fceb4
1 parent ed8e86f
commit 53fceb4
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/structuregraph_helpers/create.py b/src/structuregraph_helpers/create.py
@@ -14,6 +14,7 @@
     NearNeighbors,
     VoronoiNN,
 )
+from pymatgen.core import Structure
 
 _THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 
@@ -49,6 +50,8 @@
     "ATRCutoffDictNN",
     "LICutoffDictNN",
     "get_local_env_method",
+    "get_structure_graph",
+    "construct_clean_graph",
 )
 
 
@@ -93,6 +96,11 @@ def get_local_env_method(method: str) -> NearNeighbors:
     return VoronoiNN()
 
 
+def get_structure_graph(structure: Structure, method: str = "vesta") -> StructureGraph:
+    """Get a structure graph for a structure."""
+    return StructureGraph.with_local_env_strategy(structure, get_local_env_method(method))
+
+
 def get_nx_graph_from_edge_tuples(edge_tuples: Iterable[Tuple[int, int]]) -> nx.Graph:
     """Create a undirected graph from a list of edge tuples.