From e1f49eaabd16f9056478856b9a2af3a4c0c7d7d9 Mon Sep 17 00:00:00 2001
From: BowenD-UCB <84425382+BowenD-UCB@users.noreply.github.com>
Date: Wed, 1 Mar 2023 00:35:34 -0800
Subject: [PATCH] cleaned codes

---
 __init__.py             |  0
 examples/make_graphs.py | 36 ++++++++++++++++++++++++++----------
 2 files changed, 26 insertions(+), 10 deletions(-)
 delete mode 100644 __init__.py

diff --git a/__init__.py b/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/examples/make_graphs.py b/examples/make_graphs.py
index 5e2a9a6b..54342bfa 100644
--- a/examples/make_graphs.py
+++ b/examples/make_graphs.py
@@ -9,7 +9,10 @@
 datatype = torch.float32
 random.seed(100)
 
-# This runnable script show
+
+# This runnable script shows an example to convert a Structure json dataset to graphs
+# and save them. So the you don't have to do graph conversion in each training
+
 def main():
     data_path = ""
     graph_dir = ""
@@ -18,14 +21,22 @@ def main():
     make_graphs(data, graph_dir)
 
 
-def make_graphs(data: StructureJsonData, graph_dir, train_ratio=0.8, val_ratio=0.1):
+def make_graphs(
+        data: StructureJsonData,
+        graph_dir: str,
+        train_ratio: float = 0.8,
+        val_ratio: float = 0.1
+):
     """
-    Make cifs form the MPtrj dataset
-    :param train_ratio:
-    :param val_ratio:
-    :param test_ratio:
-    :return:
+    Make graphs from a StructureJsonData dataset
+
+    Args:
+        data (StructureJsonData): a StructureJsonData
+        graph_dir (str): a directory to save the graphs
+        train_ratio (float): train ratio
+        val_ratio (float): val ratio
     """
+    utils.mkdir(graph_dir)
     random.shuffle(data.keys)
     labels = {}
     failed_graphs = []
@@ -42,13 +53,16 @@ def make_graphs(data: StructureJsonData, graph_dir, train_ratio=0.8, val_ratio=0
             failed_graphs += [(mp_id, graph_id)]
         if i % 1000 == 0:
             print(i)
-    # torch.save(graphs, 'test_MPtrj_graphs.pt')
+
     utils.write_json(labels, os.path.join(graph_dir, "labels.json"))
     utils.write_json(failed_graphs, os.path.join(graph_dir, "failed_graphs.json"))
     make_partition(labels, graph_dir, train_ratio, val_ratio)
 
 
 def make_one_graph(mp_id, graph_id, data, graph_dir):
+    """
+    convert a structure to a Crystal_Graph and save it
+    """
     dic = data.data[mp_id].pop(graph_id)
     struc = Structure.from_dict(dic.pop("structure"))
     try:
@@ -60,9 +74,11 @@ def make_one_graph(mp_id, graph_id, data, graph_dir):
 
 
 def make_partition(
-    data, graph_dir, train_ratio=0.8, val_ratio=0.1, partition_with_frame=False
+        data, graph_dir, train_ratio=0.8, val_ratio=0.1, partition_with_frame=False
 ):
-
+    """
+    Make a train val test partition
+    """
     random.seed(42)
     if partition_with_frame is False:
         material_ids = list(data.keys())