Update graph_decoder/diffusion_utils.py

graph_decoder/diffusion_utils.py

@@ -1,131 +1,131 @@
-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):
-        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}")
-
-    with open(config_path, "r") as file:
-        cfg_dict = yaml.safe_load(file)
-
-    cfg = dict_to_namespace(cfg_dict)
-
-    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
-
-
-def to_dense(x, edge_index, edge_attr, batch, max_num_nodes=None):
-    X, node_mask = to_dense_batch(x=x, batch=batch, max_num_nodes=max_num_nodes)
-    # node_mask = node_mask.float()
-    edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
-    if max_num_nodes is None:
-        max_num_nodes = X.size(1)
-    E = to_dense_adj(
-        edge_index=edge_index,
-        batch=batch,
-        edge_attr=edge_attr,
-        max_num_nodes=max_num_nodes,
-    )
-    E = encode_no_edge(E)
-    return PlaceHolder(X=X, E=E, y=None), node_mask
-
-
-def encode_no_edge(E):
-    assert len(E.shape) == 4
-    if E.shape[-1] == 0:
-        return E
-    no_edge = torch.sum(E, dim=3) == 0
-    first_elt = E[:, :, :, 0]
-    first_elt[no_edge] = 1
-    E[:, :, :, 0] = first_elt
-    diag = (
-        torch.eye(E.shape[1], dtype=torch.bool).unsqueeze(0).expand(E.shape[0], -1, -1)
-    )
-    E[diag] = 0
-    return E
+# 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):
+#         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}")
+
+#     with open(config_path, "r") as file:
+#         cfg_dict = yaml.safe_load(file)
+
+#     cfg = dict_to_namespace(cfg_dict)
+
+#     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
+
+
+# def to_dense(x, edge_index, edge_attr, batch, max_num_nodes=None):
+#     X, node_mask = to_dense_batch(x=x, batch=batch, max_num_nodes=max_num_nodes)
+#     # node_mask = node_mask.float()
+#     edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
+#     if max_num_nodes is None:
+#         max_num_nodes = X.size(1)
+#     E = to_dense_adj(
+#         edge_index=edge_index,
+#         batch=batch,
+#         edge_attr=edge_attr,
+#         max_num_nodes=max_num_nodes,
+#     )
+#     E = encode_no_edge(E)
+#     return PlaceHolder(X=X, E=E, y=None), node_mask
+
+
+# def encode_no_edge(E):
+#     assert len(E.shape) == 4
+#     if E.shape[-1] == 0:
+#         return E
+#     no_edge = torch.sum(E, dim=3) == 0
+#     first_elt = E[:, :, :, 0]
+#     first_elt[no_edge] = 1
+#     E[:, :, :, 0] = first_elt
+#     diag = (
+#         torch.eye(E.shape[1], dtype=torch.bool).unsqueeze(0).expand(E.shape[0], -1, -1)
+#     )
+#     E[diag] = 0
+#     return E
 
 
 # #### diffusion utils