Added support for Apple silicon

.gitignore

@@ -6,3 +6,6 @@ notebooks/
 *DS_Store
 samples/
 crossdocked/
+/checkpoints
+DrugFlow on Apple Silicon.md
+/output

README.md

@@ -24,6 +24,16 @@ and add the Gnina executable for docking score computation
 wget https://github.com/gnina/gnina/releases/download/v1.1/gnina -O $CONDA_PREFIX/bin/gnina
 chmod +x $CONDA_PREFIX/bin/gnina
 ```
+### Conda Environment for Mac (Apple Silicon)
+
+Create a conda/mamba environment 
+```
+conda env create -f environment_Mac.yaml -n drugflow
+conda activate drugflow
+```
+
+Gnina is not available for Apple Silicon as it depends on libmolgrid, which heavily depends on cuda kernels. Qvina is installed in the conda environment as an alternative for docking scoring.
+
 
 ### Docker Container
 

environment_Mac.yaml

@@ -0,0 +1,30 @@
+name: drugflow2
+
+channels:
+  - conda-forge
+  - pyg
+  - bioconda
+
+dependencies:
+  - python=3.11.11
+  - pytorch=2.4.1
+  - pytorch-lightning=2.5.0
+  - rdkit=2024.03.5
+  - openbabel=3.1.1
+  - biopython=1.85
+  - scipy=1.15.2
+  - pyg=2.6.1
+  - pytorch-scatter=2.1.2
+  - ProDy=2.4.1
+  - wandb=0.16.6
+  - pandas=2.2.2
+  - posebusters=0.3.1
+  - fcd=1.2.2
+  - webdataset=0.2.100
+  - prolif=2.0.3
+  - reduce=4.14
+  - qvina=2.1.0
+  - pip=25.0.1
+  - pip:
+      - useful_rdkit_utils==0.65
+

src/data/so3_utils.py

@@ -344,7 +344,7 @@ def log_not_from_identity(point, base_point):
 
     # decorator
     if torch.__version__ >= '2.0.0':
-        GEOMSTATS_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+        GEOMSTATS_DEVICE = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
 
         def geomstats_tensor_type(func):
             def inner(*args, **kwargs):

src/generate.py

@@ -4,6 +4,7 @@
 import warnings
 import tempfile
 import pandas as pd
+import torch
 
 from Bio.PDB import PDBParser
 from pathlib import Path
@@ -66,7 +67,7 @@ def aggregate_metrics(table):
     p.add_argument('--batch_size', type=int, required=False, default=32, help="Batch size.")
     p.add_argument('--pocket_distance_cutoff', type=float, required=False, default=8.0, help="Distance cutoff to define the pocket around the reference ligand.")
     p.add_argument('--n_steps', type=int, required=False, default=None, help="Number of denoising steps.")
-    p.add_argument('--device', type=str, required=False, default='cuda:0', help="Device to use.")
+    p.add_argument('--device', type=str, required=False, default='cuda:0', help="Device to use (mps, cuda:0).")
     p.add_argument('--datadir', type=Path, required=False, default=Path(basedir, 'src', 'default'), help="Needs to be specified to sample molecule sizes.")
     p.add_argument('--seed', type=int, required=False, default=42, help="Random seed.")
     p.add_argument('--filter', action='store_true', required=False, default=False, help="Apply basic filters and keep sampling until `n_samples` molecules passing these filters are found.")
@@ -86,11 +87,19 @@ def aggregate_metrics(table):
 
     if not args.filter:
         args.batch_size = min(args.batch_size, args.n_samples)
+
+    # Select mps if available (Apple Silicon)
+    if torch.backends.mps.is_available():
+        args.device = 'mps'
 
     # Loading model
     chkpt_path = Path(args.checkpoint)
     chkpt_name = chkpt_path.parts[-1].split('.')[0]
-    model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location=args.device, strict=False)
+    if torch.backends.mps.is_available():
+        model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location='cpu', strict=False) # load model to cpu
+        model = model.to(args.device, dtype=torch.float32) # move model to mps with float32 type
+    else:
+        model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location=args.device, strict=False)
     if args.datadir is not None:
         model.datadir = args.datadir
 

src/model/gvp_transformer.py

@@ -454,7 +454,7 @@ def test_equivariance(model, nodes, edges, glob_feat):
     db = 10
     attn_heads = 9
 
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device = torch.device('mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu')
 
 
     nodes = randn(n_nodes, node_dim, device=device)