Added data and models directories

Author: igashov

.gitignore

@@ -1,4 +1,2 @@
-data
-logs
-models
-__pycache__
+.idea
+__pycache__

src/data/__init__.py

@@ -0,0 +1,113 @@
+import numpy as np
+import torch
+import pytorch_lightning as pl
+
+from torch_geometric.utils import scatter
+
+from pdb import set_trace
+
+
+class AbstractDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size, num_workers, shuffle):
+        super().__init__()
+        self.dataloaders = None
+        self.input_dims = None
+        self.output_dims = None
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.shuffle = shuffle
+
+    def train_dataloader(self):
+        return self.dataloaders["train"]
+
+    def val_dataloader(self):
+        return self.dataloaders["val"]
+
+    def test_dataloader(self):
+        return self.dataloaders["test"]
+
+    def __getitem__(self, idx):
+        return self.dataloaders['train'][idx]
+
+    def node_counts(self, max_nodes_possible=300):
+        all_counts = torch.zeros(max_nodes_possible)
+        for split in ['train', 'val', 'test']:
+            for i, data in enumerate(self.dataloaders[split]):
+                unique, counts = torch.unique(data.batch, return_counts=True)
+                for count in counts:
+                    all_counts[count] += 1
+        max_index = max(all_counts.nonzero())
+        all_counts = all_counts[:max_index + 1]
+        all_counts = all_counts / all_counts.sum()
+        return all_counts
+
+    def node_types(self):
+        num_classes = None
+        for data in self.dataloaders['train']:
+            num_classes = data.x.shape[1]
+            break
+
+        counts = torch.zeros(num_classes)
+
+        for i, data in enumerate(self.dataloaders['train']):
+            counts += data.x.sum(dim=0)
+
+        counts = counts / counts.sum()
+        return counts
+
+    def edge_counts(self):
+        num_classes = None
+        for data in self.dataloaders['train']:
+            num_classes = data.edge_attr.shape[1]
+            break
+
+        d = torch.zeros(num_classes, dtype=torch.float)
+
+        for i, data in enumerate(self.dataloaders['train']):
+            unique, counts = torch.unique(data.batch, return_counts=True)
+
+            all_pairs = 0
+            for count in counts:
+                all_pairs += count * (count - 1)
+
+            num_edges = data.edge_index.shape[1]
+            num_non_edges = all_pairs - num_edges
+
+            edge_types = data.edge_attr.sum(dim=0)
+            assert num_non_edges >= 0
+            d[0] += num_non_edges
+            d[1:] += edge_types[1:]
+
+        d = d / d.sum()
+        return d
+
+    def dummy_atoms_counts(self, max_n_dummy_nodes):
+        dummy_atoms = np.zeros(max_n_dummy_nodes + 1)
+        for data in self.dataloaders['train']:
+            batch_counts = scatter(data.p_x[:, -1], data.batch, reduce='sum')
+            for cnt in batch_counts.long().detach().cpu().numpy():
+                if cnt > max_n_dummy_nodes:
+                    continue
+                dummy_atoms[cnt] += 1
+
+        return torch.tensor(dummy_atoms) / dummy_atoms.sum()
+
+
+class MolecularDataModule(AbstractDataModule):
+    def valency_count(self, max_n_nodes):
+        valencies = torch.zeros(3 * max_n_nodes - 2)  # Max valency possible if everything is connected
+
+        # No bond, single bond, double bond, triple bond, aromatic bond
+        multiplier = torch.tensor([0, 1, 2, 3, 1.5])
+
+        for split in ['train', 'val', 'test']:
+            for i, data in enumerate(self.dataloaders[split]):
+                n = data.x.shape[0]
+
+                for atom in range(n):
+                    edges = data.edge_attr[data.edge_index[0] == atom]
+                    edges_total = edges.sum(dim=0)
+                    valency = (edges_total * multiplier).sum()
+                    valencies[valency.long().item()] += 1
+        valencies = valencies / valencies.sum()
+        return valencies