Update graph_decoder/diffusion_utils.py

graph_decoder/diffusion_utils.py

@@ -13,88 +13,88 @@ def dict_to_namespace(d):
         **{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.")
+# 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.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
+#         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}
+#         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):
-#         raise FileNotFoundError(f"Configuration file not found: {config_path}")
+def load_config(config_path, data_meta_info_path):
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Configuration file not found: {config_path}")
 
-#     if not os.path.exists(data_meta_info_path):
-#         raise FileNotFoundError(f"Data meta info file not found: {data_meta_info_path}")
+    if not os.path.exists(data_meta_info_path):
+        raise FileNotFoundError(f"Data meta info file not found: {data_meta_info_path}")
 
-#     with open(config_path, "r") as file:
-#         cfg_dict = yaml.safe_load(file)
+    with open(config_path, "r") as file:
+        cfg_dict = yaml.safe_load(file)
 
-#     cfg = dict_to_namespace(cfg_dict)
+    cfg = dict_to_namespace(cfg_dict)
 
-#     data_info = DataInfos(data_meta_info_path)
-#     return cfg, data_info
+    data_info = DataInfos(data_meta_info_path)
+    return cfg, data_info
 
 
 # #### graph utils
-# class PlaceHolder:
-#     def __init__(self, X, E, y):
-#         self.X = X
-#         self.E = E
-#         self.y = y
-
-#     def type_as(self, x: torch.Tensor, categorical: bool = False):
-#         """Changes the device and dtype of X, E, y."""
-#         self.X = self.X.type_as(x)
-#         self.E = self.E.type_as(x)
-#         if categorical:
-#             self.y = self.y.type_as(x)
-#         return self
-
-#     def mask(self, node_mask, collapse=False):
-#         x_mask = node_mask.unsqueeze(-1)  # bs, n, 1
-#         e_mask1 = x_mask.unsqueeze(2)  # bs, n, 1, 1
-#         e_mask2 = x_mask.unsqueeze(1)  # bs, 1, n, 1
-
-#         if collapse:
-#             self.X = torch.argmax(self.X, dim=-1)
-#             self.E = torch.argmax(self.E, dim=-1)
-
-#             self.X[node_mask == 0] = -1
-#             self.E[(e_mask1 * e_mask2).squeeze(-1) == 0] = -1
-#         else:
-#             self.X = self.X * x_mask
-#             self.E = self.E * e_mask1 * e_mask2
-#             assert torch.allclose(self.E, torch.transpose(self.E, 1, 2))
-#         return self
+class PlaceHolder:
+    def __init__(self, X, E, y):
+        self.X = X
+        self.E = E
+        self.y = y
+
+    def type_as(self, x: torch.Tensor, categorical: bool = False):
+        """Changes the device and dtype of X, E, y."""
+        self.X = self.X.type_as(x)
+        self.E = self.E.type_as(x)
+        if categorical:
+            self.y = self.y.type_as(x)
+        return self
+
+    def mask(self, node_mask, collapse=False):
+        x_mask = node_mask.unsqueeze(-1)  # bs, n, 1
+        e_mask1 = x_mask.unsqueeze(2)  # bs, n, 1, 1
+        e_mask2 = x_mask.unsqueeze(1)  # bs, 1, n, 1
+
+        if collapse:
+            self.X = torch.argmax(self.X, dim=-1)
+            self.E = torch.argmax(self.E, dim=-1)
+
+            self.X[node_mask == 0] = -1
+            self.E[(e_mask1 * e_mask2).squeeze(-1) == 0] = -1
+        else:
+            self.X = self.X * x_mask
+            self.E = self.E * e_mask1 * e_mask2
+            assert torch.allclose(self.E, torch.transpose(self.E, 1, 2))
+        return self
 
 
 # def to_dense(x, edge_index, edge_attr, batch, max_num_nodes=None):