add new dynamic dataset

Author: Oceanusity

OpenDFT/QHBench/QH9/config/config.yaml

@@ -1,13 +1,12 @@
 defaults:
-  - datasets: QH9-stable
+  - datasets: QH9-dynamic
 ex_name: 'ex_name'
-device: 0
+device: 2
 ckpt_dir: 'checkpoints'
 split_seed: 42
 optimizer: adam
 ema_start_epoch: -1
 
+trained_model: '.'
 # For evaluating trained model
-trained_model: './' # path to your trained model
-
 seed: 0

OpenDFT/QHBench/QH9/config/datasets/QH9-dynamic.yaml

@@ -1,5 +1,5 @@
 dataset_name: QH9Dynamic
-split: geometry  #[geometry, mol]
+split: mol  #[geometry, mol]
 
 train_batch_size: 32
 valid_batch_size: 32
@@ -13,7 +13,7 @@ pin_memory: True
 num_workers: 8
 
 warmup_steps: 1000
-total_steps: 250000
+total_steps: 260000
 lr_end: 1e-7
 
 train_batch_interval: 100

OpenDFT/QHBench/QH9/datasets.py

@@ -4,46 +4,45 @@
 import logging
 
 from models import QHNet
-from torchvision.transforms import Compose
 from torch_geometric.loader import DataLoader
+from torch_scatter import scatter_sum
 
 from datasets import QH9Stable, QH9Dynamic
 from torch_ema import ExponentialMovingAverage
 from transformers import get_polynomial_decay_schedule_with_warmup
 logger = logging.getLogger()
 
 
-def recorder_pos(data):
-    data.pos = torch.stack(
-        [data.pos[:, 2], data.pos[:, 0], data.pos[:, 1]],
-        dim=1
-    )
-    return data
-
-
 def criterion(outputs, target, loss_weights):
     error_dict = {}
     keys = loss_weights.keys()
-    # the diagonal and non-diagonal should be considered with the mask
     try:
         for key in keys:
+            row = target.edge_index[0]
+            edge_batch = target.batch[row]
             diff_diagonal = outputs[f'{key}_diagonal_blocks']-target[f'diagonal_{key}']
-            mse_diagonal  = torch.sum(diff_diagonal**2 * target[f"diagonal_{key}_mask"])
-            mae_diagonal  = torch.sum(torch.abs(diff_diagonal) * target[f"diagonal_{key}_mask"])
-            count_sum_diagonal =  torch.sum(target[f"diagonal_{key}_mask"])
+            mse_diagonal  = torch.sum(diff_diagonal**2 * target[f"diagonal_{key}_mask"], dim=[1, 2])
+            mae_diagonal  = torch.sum(torch.abs(diff_diagonal) * target[f"diagonal_{key}_mask"], dim=[1, 2])
+            count_sum_diagonal =  torch.sum(target[f"diagonal_{key}_mask"], dim=[1, 2])
+            mse_diagonal = scatter_sum(mse_diagonal, target.batch)
+            mae_diagonal = scatter_sum(mae_diagonal, target.batch)
+            count_sum_diagonal = scatter_sum(count_sum_diagonal, target.batch)
 
             diff_non_diagonal = outputs[f'{key}_non_diagonal_blocks']-target[f'non_diagonal_{key}']
-            mse_non_diagonal  = torch.sum(diff_non_diagonal**2 * target[f"non_diagonal_{key}_mask"])
-            mae_non_diagonal  = torch.sum(torch.abs(diff_non_diagonal) * target[f"non_diagonal_{key}_mask"])
-            count_sum_non_diagonal =  torch.sum(target[f"non_diagonal_{key}_mask"])
+            mse_non_diagonal  = torch.sum(diff_non_diagonal**2 * target[f"non_diagonal_{key}_mask"], dim=[1, 2])
+            mae_non_diagonal  = torch.sum(torch.abs(diff_non_diagonal) * target[f"non_diagonal_{key}_mask"], dim=[1, 2])
+            count_sum_non_diagonal =  torch.sum(target[f"non_diagonal_{key}_mask"], dim=[1, 2])
+            mse_non_diagonal = scatter_sum(mse_non_diagonal, edge_batch)
+            mae_non_diagonal = scatter_sum(mae_non_diagonal, edge_batch)
+            count_sum_non_diagonal = scatter_sum(count_sum_non_diagonal, edge_batch)
 
-            mae = (mae_diagonal + mae_non_diagonal) / (count_sum_diagonal + count_sum_non_diagonal)
-            mse = (mse_diagonal + mse_non_diagonal) / (count_sum_diagonal + count_sum_non_diagonal)
+            mae = ((mae_diagonal + mae_non_diagonal) / (count_sum_diagonal + count_sum_non_diagonal)).mean()
+            mse = ((mse_diagonal + mse_non_diagonal) / (count_sum_diagonal + count_sum_non_diagonal)).mean()
 
             error_dict[key+'_mae']  = mae
             error_dict[key+'_rmse'] = torch.sqrt(mse)
-            error_dict[key + '_diagonal_mae'] = mae_diagonal / count_sum_diagonal
-            error_dict[key + '_non_diagonal_mae'] = mae_non_diagonal / count_sum_non_diagonal
+            error_dict[key + '_diagonal_mae'] = (mae_diagonal / count_sum_diagonal).mean()
+            error_dict[key + '_non_diagonal_mae'] = (mae_non_diagonal / count_sum_non_diagonal).mean()
             loss = mse + mae
             error_dict[key] = loss
             if 'loss' in error_dict.keys():
@@ -100,6 +99,7 @@ def main(conf):
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(0)
 
+    # root_path = '/data/meng/QC_features'
     root_path = os.path.join(os.sep.join(os.getcwd().split(os.sep)[:-3]))
     # determine whether GPU is used for training
     if torch.cuda.is_available():
@@ -130,7 +130,6 @@ def main(conf):
         test_dataset, batch_size=conf.datasets.test_batch_size, shuffle=False,
         num_workers=conf.datasets.num_workers, pin_memory=conf.datasets.pin_memory)
     train_iterator = iter(train_data_loader)
-
     # define model
     model = QHNet(
         in_node_features=1,
@@ -237,7 +236,7 @@ def main(conf):
 
 def post_processing(batch, default_type):
     for key in batch.keys:
-        if torch.is_floating_point(batch[key]):
+        if torch.is_tensor(batch[key]) and torch.is_floating_point(batch[key]):
             batch[key] = batch[key].type(default_type)
     return batch
 

OpenDFT/QHBench/QH9/main.py

@@ -611,7 +611,9 @@ def forward(self, data, keep_blocks=True):
         data.node_attr, data.edge_index, data.edge_attr, data.edge_sh = \
             node_attr, edge_index, rbf_new, edge_sh
 
-        _, full_edge_index, full_edge_attr, full_edge_sh, transpose_edge_index = self.build_graph(data, 10000)
+        _, full_edge_index, full_edge_attr, full_edge_sh, transpose_edge_index = \
+            self.build_graph(data, max_radius=10000)
+
         data.full_edge_index, data.full_edge_attr, data.full_edge_sh = \
             full_edge_index, full_edge_attr, full_edge_sh
 
@@ -652,11 +654,12 @@ def forward(self, data, keep_blocks=True):
 
         return results
 
-    def build_graph(self, data, max_radius):
+    def build_graph(self, data, max_radius, edge_index=None):
         node_attr = data.atoms.squeeze()
-        
-        
-        radius_edges = radius_graph(data.pos, max_radius, data.batch, max_num_neighbors=data.num_nodes)
+        if edge_index is None:
+            radius_edges = radius_graph(data.pos, max_radius, data.batch, max_num_neighbors=data.num_nodes)
+        else:
+            radius_edges = data.full_edge_index
 
         dst, src = radius_edges
         edge_vec = data.pos[dst.long()] - data.pos[src.long()]

OpenDFT/QHBench/QH9/models/QHNet.py

@@ -1,2 +1,3 @@
-from .QHNet import QHNet
+# from .QHNet import QHNet
+from .ori_QHNet_with_bias import QHNet