Initial incorporation of raw coordinate training

bin/train.py

@@ -49,7 +49,7 @@ torch.backends.cudnn.benchmark = False
 
 # Define some typing literals
 ANGLES_DEFINITIONS = Literal[
-    "canonical", "canonical-full-angles", "canonical-minimal-angles"
+    "canonical", "canonical-full-angles", "canonical-minimal-angles", "cart-coords"
 ]
 
 
@@ -141,6 +141,7 @@ def get_train_valid_test_sets(
         "canonical": datasets.CathCanonicalAnglesDataset,
         "canonical-full-angles": datasets.CathCanonicalAnglesOnlyDataset,
         "canonical-minimal-angles": datasets.CathCanonicalMinimalAnglesDataset,
+        "cart-coords": datasets.CathCanonicalCoordsDataset,
     }[angles_definitions]
     logging.info(f"Clean dataset class: {clean_dset_class}")
 
@@ -152,7 +153,7 @@ def get_train_valid_test_sets(
             pad=max_seq_len,
             min_length=min_seq_len,
             trim_strategy=seq_trim_strategy,
-            zero_center=True,
+            zero_center=False if angles_definitions == "cart-coords" else True,
             toy=toy,
         )
         for s in splits

config_jsons/cath_full_angles_cosine.json

@@ -1,6 +1,5 @@
 {
     "angles_definitions": "canonical-full-angles",
-    "zero_center": true,
     "timesteps": 1000,
     "max_seq_len": 128,
     "min_seq_len": 40,

config_jsons/cath_raw_coordinates.json

@@ -0,0 +1,30 @@
+{
+    "angles_definitions": "cart-coords",
+    "timesteps": 1000,
+    "max_seq_len": 128,
+    "min_seq_len": 40,
+    "trim_strategy": "randomcrop",
+    "variance_schedule": "cosine",
+    "variance_scale": 1.0,
+    "time_encoding": "gaussian_fourier",
+    "position_embedding_type": "relative_key",
+    "num_hidden_layers": 12,
+    "hidden_size": 384,
+    "intermediate_size": 768,
+    "num_heads": 12,
+    "dropout_p": 0.1,
+    "decoder": "mlp",
+    "gradient_clip": 1.0,
+    "lr": 5e-5,
+    "loss": "smooth_l1",
+    "l2_norm": 0.0,
+    "l1_norm": 0.0,
+    "circle_reg": 0.0,
+    "lr_scheduler": "LinearWarmup",
+    "min_epochs": 10000,
+    "max_epochs": 10000,
+    "early_stop_patience": 0,
+    "use_swa": false,
+    "batch_size": 64,
+    "multithread": true
+}