Linting and updates

.pre-commit-config.yaml

@@ -0,0 +1,25 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.5.0
+  hooks:
+    - id: trailing-whitespace
+    - id: end-of-file-fixer
+    - id: check-yaml
+    - id: check-added-large-files
+- repo: https://github.com/ambv/black
+  rev: 23.9.1
+  hooks:
+    - id: black
+- repo: https://github.com/jsh9/pydoclint
+  # pydoclint version.
+  rev: 0.3.3
+  hooks:
+    - id: pydoclint
+      args:
+        - "--config=pyproject.toml"
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.1.1
+  hooks:
+    - id: ruff
+      args: [--fix, --exit-non-zero-on-fix]

docs/source/conf.py

@@ -109,7 +109,6 @@
             "vu": "\\mathbf{u}",
             "vv": "\\mathbf{v}",
             "vw": "\\mathbf{w}",
-            "vx": "\\mathbf{x}",
             "vy": "\\mathbf{y}",
             "vz": "\\mathbf{z}",
         }

proteinworkshop/config/dataset/cath.yaml

@@ -9,4 +9,5 @@ datamodule:
   dataset_fraction: 1.0 # Fraction of the dataset to use
   transforms: ${transforms} # Transforms to apply to dataset examples
   overwrite: False # Whether to overwrite the dataset if it already exists
+  in_memory: True # Whether to load the entire dataset into memory
 num_classes: 23 # Number of classes

proteinworkshop/config/dataset/fold_family.yaml

@@ -9,4 +9,5 @@ datamodule:
   shuffle_labels: False # Whether to shuffle labels for permutation testing
   transforms: ${transforms} # Transforms to apply to dataset examples
   overwrite: False # Whether to overwrite existing dataset files
+  in_memory: True # Whether to load the entire dataset into memory
 num_classes: 1195 # Number of classes

proteinworkshop/config/dataset/fold_fold.yaml

@@ -9,4 +9,5 @@ datamodule:
   shuffle_labels: False # Whether to shuffle labels for permutation testing
   transforms: ${transforms} # Transforms to apply to dataset examples
   overwrite: False # Whether to overwrite existing dataset files
+  in_memory: True # Whether to load the entire dataset into memory
 num_classes: 1195 # Number of classes

proteinworkshop/config/dataset/fold_superfamily.yaml

@@ -9,4 +9,5 @@ datamodule:
   shuffle_labels: False # Whether to shuffle labels for permutation testing
   transforms: ${transforms} # Transforms to apply to dataset examples
   overwrite: False # Whether to overwrite existing dataset files
+  in_memory: True # Whether to load the entire dataset into memory
 num_classes: 1195 # Number of classes

proteinworkshop/config/finetune.yaml

@@ -25,6 +25,7 @@ defaults:
   - finetune: default # Specifies finetuning config. See: proteinworkshop/config/finetune/
   # debugging config (enable through command line, e.g. `python train.py debug=default)
   - debug: null
+  - optional hparams: ${encoder}_${features}
   - _self_ # see: https://hydra.cc/docs/upgrades/1.0_to_1.1/default_composition_order/. Adding _self_ at bottom means values in this file override defaults.
 
 task_name: "finetune"

proteinworkshop/config/hparams/egnn_ca_angles.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/egnn_ca_base.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/egnn_ca_bb.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/egnn_ca_sc.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/egnn_ca_seq.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/gcpnet_ca_angles.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/gcpnet_ca_base.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/gcpnet_ca_bb.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/gcpnet_ca_sc.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/gcpnet_ca_seq.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/gear_net_edge_ca_angles.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/gear_net_edge_ca_base.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/gear_net_edge_ca_bb.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/gear_net_edge_ca_sc.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.1

proteinworkshop/config/hparams/gear_net_edge_ca_seq.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/mace_ca_base.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0010
+  decoder_dropout: 0.0

proteinworkshop/config/hparams/mace_ca_bb.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/mace_ca_sc.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/mace_ca_seq.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/schnet_ca_angles.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.1

proteinworkshop/config/hparams/schnet_ca_base.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/schnet_ca_bb.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/schnet_ca_sc.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/schnet_ca_seq.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/tfn_ca_angles.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/tfn_ca_base.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/tfn_ca_bb.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.001
+  decoder_dropout: 0.5

proteinworkshop/config/hparams/tfn_ca_sc.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0003
+  decoder_dropout: 0.3

proteinworkshop/config/hparams/tfn_ca_seq.yaml

@@ -0,0 +1,3 @@
+hparams:
+  lr: 0.0001
+  decoder_dropout: 0.3

proteinworkshop/config/sweeps/baseline_inverse_folding.yaml

@@ -7,19 +7,16 @@ metric: # Does not matter, as we are using sweep to run the experiment.
 
 parameters:
   task:
-    values: [inverse_folding]
+    value: inverse_folding
 
   dataset:
-    values: [cath]
+    value: cath
 
   encoder:
-    values: [schnet, dimenet_plus_plus, egnn, gcpnet, gear_net_edge]
-
-  optimiser.optimizer.lr:
-    values: [0.0001]
+    values: [schnet, egnn, gcpnet, gear_net_edge, tfn, mace]
 
   features:
-    values: [ca_base, ca_seq, ca_angles, ca_bb]
+    values: [ca_seq, ca_angles, ca_bb]
 
   scheduler:
     value: plateau
@@ -28,14 +25,32 @@ parameters:
     value: False
 
   +aux_task:
-    values: [none, nn_sequence, nn_structure_torsion, nn_structure_r3]
+    values: [none, nn_structure_torsion, nn_structure_r3]
 
   trainer.max_epochs:
-    value: 150
+    value: 250
 
   test:
     value: True
 
+  trainer:
+    value: gpu
+
+  logger:
+    value: wandb
+
+  seed:
+    values: [13, 42, 121]
+
+  name:
+    value: "${hydra:runtime.choices.encoder}_${hydra:runtime.choices.features}_${hydra:runtime.choices.aux_task}_seed_${seed}"
+
+  optimiser.optimizer.lr:
+    value: ${hparams.hparams.lr}
+
+  decoder.residue_type.dropout:
+    value: ${hparams.hparams.decoder_dropout}
+
 command:
   - ${env}
   - HYDRA_FULL_ERROR=1

proteinworkshop/config/sweeps/find_hparams.yaml

@@ -0,0 +1,54 @@
+program: proteinworkshop/train.py
+method: grid
+name: baseline_hyperparameter_search
+metric: # Does not matter, as we are using sweep to run the experiment.
+  goal: minimize
+  name: val/loss/total
+
+parameters:
+  task:
+    values: [multiclass_graph_classification]
+
+  dataset:
+    values: [fold_family]
+
+  encoder:
+    values: [schnet, gear_net_edge, egnn, gcpnet] #, tfn]
+
+  optimiser.optimizer.lr:
+    values: [0.00001, 0.0001, 0.0003, 0.001]
+
+  decoder.graph_label.dropout:
+    values: [0.0, 0.1, 0.3, 0.5]
+
+  features:
+    values: [ca_base, ca_seq, ca_angles, ca_bb, ca_sc]
+
+  scheduler:
+    value: plateau
+
+  extras.enforce_tags:
+    value: False
+
+  #+aux_task:
+  #  values: [none, nn_sequence, nn_structure_torsion, nn_structure_r3]
+
+  trainer.max_epochs:
+    value: 300
+
+  test:
+    value: True
+
+  logger:
+    value: wandb
+
+  name:
+    value: "${hydra:runtime.choices.encoder}_${hydra:runtime.choices.features}_lr_${optimiser.optimizer.lr}_d_${decoder.graph_label.dropout}"
+
+command:
+  - ${env}
+  - HYDRA_FULL_ERROR=1
+  - WANDB_START_METHOD=thread
+  - python
+  - ${program}
+  - ${args_no_hyphens}

proteinworkshop/config/sweeps/pre_train_egnn.yaml

@@ -0,0 +1,55 @@
+program: proteinworkshop/train.py
+method: grid
+name: pretrain_egnn
+metric: # Does not matter, as we are using sweep to run the experiment.
+  goal: minimize
+  name: val/loss/total
+
+parameters:
+  task:
+    values:
+      [
+        inverse_folding,
+        sequence_denoising,
+        #plddt_prediction,
+        structure_denoising,
+        torsional_denoising,
+      ]
+
+  dataset:
+    value: afdb_rep_v4
+
+  dataset.datamodule.num_workers:
+    value: 16
+
+  encoder:
+    values: [egnn]
+
+  optimiser.optimizer.lr:
+    values: [0.0001]
+
+  features:
+    values: [ca_angles, ca_bb]
+
+  scheduler:
+    value: linear_warmup_cosine_decay
+
+  callbacks.model_checkpoint.every_n_epochs:
+    value: 1
+
+  extras.enforce_tags:
+    value: False
+
+  trainer:
+    value: ddp
+
+  +trainer.max_epochs:
+    value: 10
+
+command:
+  - ${env}
+  - HYDRA_FULL_ERROR=1
+  - WANDB_START_METHOD=thread
+  - python
+  - ${program}
+  - ${args_no_hyphens}