Minor bug fixes

Author: Old-Shatterhand

configs/lgi/all_cold_glycan.yaml

@@ -0,0 +1,47 @@
+seed: 42
+# root_dir: /scratch/SCRATCH_SAS/roman/Gothenburg/GIFFLAR/data_new_256
+# logs_dir: /scratch/SCRATCH_SAS/roman/Gothenburg/GIFFLAR/logs_new_256
+# origin: /scratch/SCRATCH_SAS/roman/Gothenburg/GIFFLAR/data_new_256/lgi_random_small.csv
+root_dir: /scratch/data_new_256
+logs_dir: /scratch/logs_new_256
+origin: /scratch/data_new_256/lgi_cold_glycan_150k.csv
+pre-transforms:
+  GIFFLARTransform:
+  GNNGLYTransform:
+  ECFPTransform:
+  SweetNetTransform:
+model:
+  glycan_encoder:
+    - name: mlp
+      feat_dim: 1024
+      hidden_dim: 1024
+      output_dim: 1024
+      num_layers: 1
+    - name: sweetnet
+      feat_dim: 128
+      hidden_dim: 256
+      num_layers: 16
+      suffix:
+    - name: gnngly
+      feat_dim: 128
+      hidden_dim: 256
+      num_layers: 8
+      suffix:
+    # - name: rgcn
+    #   feat_dim: 128
+    #   hidden_dim: 256
+    #   num_layers: 8
+    #   suffix:
+    - name: gifflar
+      feat_dim: 128
+      hidden_dim: 256
+      num_layers: 8
+      pooling: global_mean
+      suffix:
+  lectin_encoder:
+    - name: ESM
+      layer_num: 33
+  batch_size: 256
+  epochs: 25
+  learning_rate: 0.001
+  optimizer: Adam

experiments/lectinoracle.py

@@ -7,7 +7,7 @@
 from glycowork.ml.models import prep_model
 
 from gifflar.data.modules import LGI_GDM
-from gifflar.data.datasets import GlycanOnDiskDataset
+from gifflar.data.datasets_lgi import GlycanOnDiskDataset
 from experiments.protein_encoding import ENCODER_MAP
 from gifflar.model.utils import LectinStorage
 

experiments/lgi_model.py

@@ -9,6 +9,8 @@
 from experiments.protein_encoding import ENCODER_MAP, EMBED_SIZES
 from gifflar.data.hetero import HeteroDataBatch
 from gifflar.model.base import GlycanGIN
+from gifflar.model.baselines.gnngly import GNNGLY
+from gifflar.model.baselines.mlp import MLP
 from gifflar.model.baselines.sweetnet import SweetNetLightning
 from gifflar.model.downstream import DownstreamGGIN
 from gifflar.model.glylm import GlycanLM
@@ -19,6 +21,8 @@
 
 GLYCAN_ENCODERS = {
     "gifflar": DownstreamGGIN,
+    "gnngly": GNNGLY,
+    "mlp": MLP,
     "sweetnet": SweetNetLightning,
     "glylm": GlycanLM,
 }
@@ -110,7 +114,7 @@ def shared_step(self, batch: HeteroData, stage: str, batch_idx: int = 0, dataloa
         """
         fwd_dict = self(batch)
         fwd_dict["labels"] = batch["y"]
-        fwd_dict["preds"] = fwd_dict["preds"].reshape(-1)
+        fwd_dict["preds"] = fwd_dict["preds"].reshape(-1).float()
         if dataloader_idx == 0:
             fwd_dict["loss"] = self.loss(fwd_dict["preds"], fwd_dict["labels"])
             self.metrics[stage].update(fwd_dict["preds"], fwd_dict["labels"])

experiments/train_lgi.py

@@ -54,7 +54,7 @@ def unfold_config(config: dict):
 
 
 def train(contrastive: bool = False, ckpt_file: Path | None = None, **kwargs):
-    kwargs["pre-transforms"] = {"GIFFLARTransform": "", "SweetNetTransform": ""}
+    # kwargs["pre-transforms"] = {"GIFFLARTransform": "", "SweetNetTransform": ""}
     kwargs["hash"] = hash_dict(kwargs["pre-transforms"])
     seed_everything(kwargs["seed"])
 

ext_baselines/sweettalk.py

@@ -171,6 +171,9 @@ def main(base: Path, task: str):
     elif task == "Taxonomy_Kingdom":
         config = {"name": "Taxonomy_Kingdom", "task": "classification", "num_classes": 13}
         metrics = get_metrics("multilabel", n_outputs=13)
+    elif task == "spectrum":
+        config = {"name": "Spectrum", "task": "regression", "num_classes": 2048}
+        metrics = get_metrics("regression", n_outputs=2048)
     else:
         raise ValueError(f"Unknown task {task}")
 
@@ -202,7 +205,7 @@ def main(base: Path, task: str):
         scheduler, 
         metrics,
         datamodule,
-        num_epochs=100, 
+        num_epochs=2, 
         padding=False
     )
     torch.save(model.state_dict(), version / "model.pth")
@@ -211,6 +214,6 @@ def main(base: Path, task: str):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--base", type=Path, required=True)
-    parser.add_argument("--task", type=str, required=True, choices=["glycosylation", "tissue", "Taxonomy_Kingdom"])
+    parser.add_argument("--task", type=str, required=True, choices=["glycosylation", "tissue", "Taxonomy_Kingdom", "spectrum"])
     args = parser.parse_args()
     main(args.base, args.task)