Update graph_decoder/diffusion_utils.py

graph_decoder/diffusion_utils.py

@@ -1,52 +1,52 @@
-# import os
-# import json
-# import yaml
-
-# import torch
-# import numpy as np
-# from torch.nn import functional as F
-# from torch_geometric.utils import to_dense_adj, to_dense_batch, remove_self_loops
-# from types import SimpleNamespace
-
-# def dict_to_namespace(d):
-#     return SimpleNamespace(
-#         **{k: dict_to_namespace(v) if isinstance(v, dict) else v for k, v in d.items()}
-#     )
-
-# class DataInfos:
-#     def __init__(self, meta_filename="data.meta.json"):
-#         self.all_targets = ['CH4', 'CO2', 'H2', 'N2', 'O2']
-#         self.task_type = "gas_permeability"
-#         if os.path.exists(meta_filename):
-#             with open(meta_filename, "r") as f:
-#                 meta_dict = json.load(f)
-#         else:
-#             raise FileNotFoundError(f"Meta file {meta_filename} not found.")
-
-#         self.active_atoms = meta_dict["active_atoms"]
-#         self.max_n_nodes = meta_dict["max_node"]
-#         self.original_max_n_nodes = meta_dict["max_node"]
-#         self.n_nodes = torch.Tensor(meta_dict["n_atoms_per_mol_dist"])
-#         self.edge_types = torch.Tensor(meta_dict["bond_type_dist"])
-#         self.transition_E = torch.Tensor(meta_dict["transition_E"])
-
-#         self.atom_decoder = meta_dict["active_atoms"]
-#         node_types = torch.Tensor(meta_dict["atom_type_dist"])
-#         active_index = (node_types > 0).nonzero().squeeze()
-#         self.node_types = torch.Tensor(meta_dict["atom_type_dist"])[active_index]
-#         self.nodes_dist = DistributionNodes(self.n_nodes)
-#         self.active_index = active_index
-
-#         val_len = 3 * self.original_max_n_nodes - 2
-#         meta_val = torch.Tensor(meta_dict["valencies"])
-#         self.valency_distribution = torch.zeros(val_len)
-#         val_len = min(val_len, len(meta_val))
-#         self.valency_distribution[:val_len] = meta_val[:val_len]
-#         ## for all
-#         self.input_dims = {"X": len(self.active_atoms), "E": 5, "y": 5}
-#         self.output_dims = {"X": len(self.active_atoms), "E": 5, "y": 5}
-#         # self.input_dims = {"X": 11, "E": 5, "y": 5}
-#         # self.output_dims = {"X": 11, "E": 5, "y": 5}
+import os
+import json
+import yaml
+
+import torch
+import numpy as np
+from torch.nn import functional as F
+from torch_geometric.utils import to_dense_adj, to_dense_batch, remove_self_loops
+from types import SimpleNamespace
+
+def dict_to_namespace(d):
+    return SimpleNamespace(
+        **{k: dict_to_namespace(v) if isinstance(v, dict) else v for k, v in d.items()}
+    )
+
+class DataInfos:
+    def __init__(self, meta_filename="data.meta.json"):
+        self.all_targets = ['CH4', 'CO2', 'H2', 'N2', 'O2']
+        self.task_type = "gas_permeability"
+        if os.path.exists(meta_filename):
+            with open(meta_filename, "r") as f:
+                meta_dict = json.load(f)
+        else:
+            raise FileNotFoundError(f"Meta file {meta_filename} not found.")
+
+        self.active_atoms = meta_dict["active_atoms"]
+        self.max_n_nodes = meta_dict["max_node"]
+        self.original_max_n_nodes = meta_dict["max_node"]
+        self.n_nodes = torch.Tensor(meta_dict["n_atoms_per_mol_dist"])
+        self.edge_types = torch.Tensor(meta_dict["bond_type_dist"])
+        self.transition_E = torch.Tensor(meta_dict["transition_E"])
+
+        self.atom_decoder = meta_dict["active_atoms"]
+        node_types = torch.Tensor(meta_dict["atom_type_dist"])
+        active_index = (node_types > 0).nonzero().squeeze()
+        self.node_types = torch.Tensor(meta_dict["atom_type_dist"])[active_index]
+        self.nodes_dist = DistributionNodes(self.n_nodes)
+        self.active_index = active_index
+
+        val_len = 3 * self.original_max_n_nodes - 2
+        meta_val = torch.Tensor(meta_dict["valencies"])
+        self.valency_distribution = torch.zeros(val_len)
+        val_len = min(val_len, len(meta_val))
+        self.valency_distribution[:val_len] = meta_val[:val_len]
+        ## for all
+        self.input_dims = {"X": len(self.active_atoms), "E": 5, "y": 5}
+        self.output_dims = {"X": len(self.active_atoms), "E": 5, "y": 5}
+        # self.input_dims = {"X": 11, "E": 5, "y": 5}
+        # self.output_dims = {"X": 11, "E": 5, "y": 5}
 
 # def load_config(config_path, data_meta_info_path):
 #     if not os.path.exists(config_path):