AIRI-Institute · KuzmaKhrabrov · May 13, 2024 · Mar 5, 2024 · Mar 5, 2024 · Mar 12, 2024
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -1,2 +1,4 @@
 include nablaDFT/links/*
 include nablaDFT/data/*
+include nablaDFT/equiformer_v2/Jd.pt
+include nablaDFT/escn/Jd.pt
diff --git a/config/datamodule/nablaDFT_ase.yaml b/config/datamodule/nablaDFT_ase.yaml
@@ -1,7 +1,6 @@
 # Dataset config for ASE nablaDFT
-_target_: nablaDFT.dataset.NablaDFT
+_target_: nablaDFT.dataset.ASENablaDFT
 
-type_of_nn: ASE
 split: ${job_type}
 datapath: ./datasets/nablaDFT/train/raw
 dataset_name: ${dataset_name}

diff --git a/config/datamodule/nablaDFT_ase_test.yaml b/config/datamodule/nablaDFT_ase_test.yaml
@@ -1,7 +1,6 @@
 # Dataset config for ASE nablaDFT
-_target_: nablaDFT.dataset.NablaDFT
+_target_: nablaDFT.dataset.ASENablaDFT
 
-type_of_nn: ASE
 split: ${job_type}
 datapath: ./datasets/nablaDFT/test/raw
 dataset_name: ${dataset_name}

diff --git a/config/datamodule/nablaDFT_hamiltonian.yaml b/config/datamodule/nablaDFT_hamiltonian.yaml
@@ -0,0 +1,13 @@
+# Dataset config for torch geometric nablaDFT
+_target_: nablaDFT.dataset.PyGHamiltonianDataModule
+
+root: ./datasets/nablaDFT/hamiltonian
+dataset_name: ${dataset_name}
+train_size: 0.9
+val_size: 0.1
+seed: 23
+# Dataloader args
+batch_size: 8
+num_workers: 4
+persistent_workers: True
+pin_memory: True
diff --git a/config/datamodule/nablaDFT_hamiltonian_test.yaml b/config/datamodule/nablaDFT_hamiltonian_test.yaml
@@ -0,0 +1,11 @@
+# Dataset config for torch geometric nablaDFT
+_target_: nablaDFT.dataset.PyGHamiltonianDataModule
+
+root: ./datasets/nablaDFT/hamiltonian
+dataset_name: ${dataset_name}
+seed: 23
+# Dataloader args
+batch_size: 2
+num_workers: 4
+persistent_workers: True
+pin_memory: True
diff --git a/config/datamodule/nablaDFT_pyg.yaml b/config/datamodule/nablaDFT_pyg.yaml
@@ -1,12 +1,14 @@
 # Dataset config for torch geometric nablaDFT
-_target_: nablaDFT.dataset.NablaDFT
+_target_: nablaDFT.dataset.PyGNablaDFTDataModule
 
-type_of_nn: PyG
-split: ${job_type}
 root: ./datasets/nablaDFT/train
 dataset_name: ${dataset_name}
 train_size: 0.9
 val_size: 0.1
+seed: 23
+# Dataloader args
 batch_size: 32
 num_workers: 8
-seed: 23
+persistent_workers: True
+pin_memory: True
+
diff --git a/config/datamodule/nablaDFT_pyg_test.yaml b/config/datamodule/nablaDFT_pyg_test.yaml
@@ -1,9 +1,10 @@
 # Dataset config for torch geometric nablaDFT
-_target_: nablaDFT.dataset.NablaDFT
+_target_: nablaDFT.dataset.PyGNablaDFTDataModule
 
-type_of_nn: PyG
-split: ${job_type}
 root: ./datasets/nablaDFT/test
 dataset_name: ${dataset_name}
+# Dataloader args
 batch_size: 32
-num_workers: 12
+num_workers: 12
+persistent_workers: True
+pin_memory: True
diff --git a/config/dimenetplusplus.yaml b/config/dimenetplusplus.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: DimeNet++
-dataset_name: dataset_train_2k
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 job_type: train
 pretrained: False

diff --git a/config/equiformer_v2_oc20.yaml b/config/equiformer_v2_oc20.yaml
@@ -1,6 +1,6 @@
 # Global variables
-name: Equiformer_v2_OC20
-dataset_name: dataset_train_2k
+name: Equiformer_v2
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 warmup_steps: 0
 job_type: train

diff --git a/config/escn-oc.yaml b/config/escn-oc.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: ESCN-OC
-dataset_name: dataset_train_2k
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 warmup_steps: 0
 job_type: train

diff --git a/config/gemnet-oc.yaml b/config/gemnet-oc.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: GemNet-OC
-dataset_name: dataset_train_2k
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 warmup_steps: 0
 job_type: train

diff --git a/config/gemnet-oc_test.yaml b/config/gemnet-oc_test.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: GemNet-OC
-dataset_name: dataset_test_conformations_2k
+dataset_name: dataset_test_conformations_tiny
 max_steps: 1000000
 warmup_steps: 0
 job_type: test

diff --git a/config/graphormer3d.yaml b/config/graphormer3d.yaml
@@ -1,6 +1,6 @@
 # Global variables
-name: Graphormer3D-half
-dataset_name: dataset_train_2k
+name: Graphormer3D-small
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 warmup_steps: 60000
 job_type: train

diff --git a/config/model/graphormer3d-base.yaml b/config/model/graphormer3d-base.yaml
@@ -20,7 +20,7 @@ optimizer:
   lr: 3e-4
 
 lr_scheduler:
-  _target_: nablaDFT.graphormer.schedulers.get_linear_schedule_with_warmup
+  _target_: nablaDFT.schedulers.get_linear_schedule_with_warmup
   _partial_: true
   num_warmup_steps: ${warmup_steps}
   num_training_steps: ${max_steps}

diff --git a/config/model/graphormer3d-half.yaml b/config/model/graphormer3d-half.yaml
@@ -20,7 +20,7 @@ optimizer:
   lr: 3e-4
 
 lr_scheduler:
-  _target_: nablaDFT.graphormer.schedulers.get_linear_schedule_with_warmup
+  _target_: nablaDFT.schedulers.get_linear_schedule_with_warmup
   _partial_: true
   num_warmup_steps: ${warmup_steps}
   num_training_steps: ${max_steps}

diff --git a/config/model/graphormer3d-large.yaml b/config/model/graphormer3d-large.yaml
@@ -22,7 +22,7 @@ optimizer:
   weight_decay: 1e-3
 
 lr_scheduler:
-  _target_: nablaDFT.graphormer.schedulers.get_linear_schedule_with_warmup
+  _target_: nablaDFT.schedulers.get_linear_schedule_with_warmup
   _partial_: true
   num_warmup_steps: ${warmup_steps}
   num_training_steps: ${max_steps}

diff --git a/config/model/painn-oc.yaml b/config/model/painn-oc.yaml
@@ -3,11 +3,11 @@ _target_: nablaDFT.painn_pyg.PaiNNLightning
 model_name: "PAINN-OC"
 net:
   _target_: nablaDFT.painn_pyg.PaiNN
-  hidden_channels: 512
+  hidden_channels: 128
   num_layers: 6
-  num_rbf: 128
-  cutoff: 12.0
-  max_neighbors: 50
+  num_rbf: 100
+  cutoff: 5.0
+  max_neighbors: 100
   rbf: 
       name: 'gaussian'
   envelope: 
@@ -17,22 +17,21 @@ net:
   direct_forces: false
   use_pbc: false
   otf_graph: true
-  num_elements: 65
+  num_elements: 100
 
 
 optimizer:
   _target_: torch.optim.AdamW
   _partial_: true
-  amsgrad: true
-  betas: [0.9, 0.95]
-  lr: 1e-3
+  lr: 5e-4
   weight_decay: 0
 
 lr_scheduler:
   _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
   _partial_: true
   factor: 0.8
-  patience: 10
+  patience: 100
+  min_lr: 1e-6
 
 losses:
   energy:
@@ -41,12 +40,7 @@ losses:
     _target_: nablaDFT.gemnet_oc.loss.L2Loss
 loss_coefs:
   energy: 1.0
-  forces: 100.0
-
-ema:
-  _target_: torch_ema.ExponentialMovingAverage
-  _partial_: true
-  decay: 0.9999
+  forces: 1.0
 
 metric:
   _target_: torchmetrics.MultitaskWrapper
@@ -55,4 +49,4 @@ metric:
     energy:
       _target_: torchmetrics.MeanAbsoluteError
     forces:
-      _target_: torchmetrics.MeanAbsoluteError
+      _target_: torchmetrics.MeanAbsoluteError
diff --git a/config/model/qhnet.yaml b/config/model/qhnet.yaml
@@ -0,0 +1,54 @@
+_target_: nablaDFT.qhnet.QHNetLightning
+
+model_name: "QHNet"
+net:
+  _target_: nablaDFT.qhnet.QHNet
+  _convert_: partial
+  sh_lmax: 4
+  hidden_size: 128
+  bottle_hidden_size: 32
+  num_gnn_layers: 5
+  max_radius: 12
+  num_nodes: 83
+  radius_embed_dim: 32
+  orbitals:
+    1: [0, 0, 1]
+    6: [0, 0, 0, 1, 1, 2]
+    7: [0, 0, 0, 1, 1, 2]
+    8: [0, 0, 0, 1, 1, 2]
+    9: [0, 0, 0, 1, 1, 2]
+    16: [0, 0, 0, 0, 1, 1, 1, 2]
+    17: [0, 0, 0, 0, 1, 1, 1, 2]
+    35: [0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2]
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  amsgrad: true
+  betas: [0.9, 0.95]
+  lr: 5e-4
+
+lr_scheduler:
+  _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
+  _partial_: true
+  factor: 0.8
+  patience: 10
+  min_lr: 1e-6
+
+losses:
+  hamiltonian:
+    _target_: nablaDFT.qhnet.loss.HamiltonianLoss
+loss_coefs:
+  hamiltonian: 1.0
+
+metric:
+  _target_: torchmetrics.MultitaskWrapper
+  _convert_: all
+  task_metrics:
+    hamiltonian:
+      _target_: nablaDFT.qhnet.MaskedMeanAbsoluteError
+
+ema:
+  _target_: torch_ema.ExponentialMovingAverage
+  _partial_: true
+  decay: 0.9999
diff --git a/config/painn-oc.yaml b/config/painn-oc.yaml
@@ -1,8 +1,8 @@
 # Global variables
 name: PaiNN-OC
-dataset_name: dataset_train_2k
+dataset_name: dataset_train_tiny
 max_steps: 1000000
-job_type: test
+job_type: train
 pretrained: False
 ckpt_path: null # path to checkpoint for training resume or test run
 

diff --git a/config/painn.yaml b/config/painn.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: PaiNN
-dataset_name: dataset_train_2k
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 job_type: train
 pretrained: False

diff --git a/config/qhnet.yaml b/config/qhnet.yaml
@@ -0,0 +1,26 @@
+# Global variables
+name: QHNet
+dataset_name: dataset_train_tiny
+max_steps: 1000000
+warmup_steps: 0
+job_type: train
+pretrained: False
+ckpt_path: null # path to checkpoint for training resume or test run
+
+# configs
+defaults:
+  - _self_
+  - datamodule: nablaDFT_hamiltonian.yaml  # dataset config
+  - model: qhnet.yaml  # model config
+  - callbacks: default.yaml  # pl callbacks config
+  - loggers: wandb.yaml  # pl loggers config
+  - trainer: train.yaml  # trainer config
+
+# need this to set working dir as current dir
+hydra:
+  output_subdir: null
+  run:
+    dir: .
+original_work_dir: ${hydra:runtime.cwd}
+
+seed: 23
diff --git a/config/schnet.yaml b/config/schnet.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: SchNet
-dataset_name: dataset_train_2k
+dataset_name: dataset_train_tiny
 max_steps: 1000000
 job_type: train
 pretrained: False

diff --git a/config/schnet_test.yaml b/config/schnet_test.yaml
@@ -1,6 +1,6 @@
 # Global variables
 name: SchNet
-dataset_name: dataset_test_conformations_2k
+dataset_name: dataset_test_conformations_tiny
 max_steps: 1000000
 job_type: test
 pretrained: False

diff --git a/config/trainer/train.yaml b/config/trainer/train.yaml
@@ -5,7 +5,8 @@ accelerator: "gpu"
 devices: [0]
 strategy:
   _target_: pytorch_lightning.strategies.ddp.DDPStrategy
-
+  # QHNet has unused params, uncomment line for train
+  # find_unused_parameters: True
 max_steps: ${max_steps}
 
 # example of additional arguments for trainer

diff --git a/nablaDFT/__init__.py b/nablaDFT/__init__.py
@@ -7,3 +7,5 @@
 from . import escn
 from . import ase_model
 from . import painn_pyg
+
+from . import schedulers