Exclude non-standard residues from ConsecutiveCA filter

README.md

@@ -25,7 +25,7 @@ pip install boltzgen
 Choose the installer for your operating system, download it, and follow the on-screen prompts:
 
 * **Windows:** <https://www.anaconda.com/docs/getting-started/miniconda/install#windows-installation>
-* **macOS / Linux:** <https://www.anaconda.com/docs/getting-started/miniconda/install#macos-linux-installation>
+* **MacOS / Linux:** <https://www.anaconda.com/docs/getting-started/miniconda/install#macos-linux-installation>
 
 After installation, **open a terminal / command prompt** (you may need to search for “Anaconda Prompt” on Windows).
 
@@ -36,6 +36,17 @@ Run the command below in a terminal to create a fresh environment called `bg` wi
 ```bash
 conda create -n bg python=3.12
 ```
+* **MacOS**
+
+Create a new conda environment for boltzgen with python 3.12, numba, numpy and lvmlite:
+
+```
+conda create --name bg python=3.12 llvmlite==0.44.0 numba==0.61.0 numpy==2.0.2
+```
+Temporary fix for loading multiple libomp
+```
+export KMP_DUPLICATE_LIB_OK=TRUE
+```
 
 ### 3 - Activate the environment (do this every time you work with BoltzGen)
 
@@ -97,7 +108,6 @@ docker build -t boltzgen:weights --build-arg DOWNLOAD_WEIGHTS=true .
 ⚠️ it downloads models (~6GB) to `~/.cache`. This can by changed by passing `--cache YOUR_PATH` or by setting `$HF_HOME`.\
 ⚠️ If your run is ever interrupted, you can restart it with `--reuse`. No progress is lost.
 
-
 ```bash
 boltzgen run example/vanilla_protein/1g13prot.yaml \
   --output workbench/test_run \

pyproject.toml

@@ -14,8 +14,8 @@ readme = { file = "PYPI_DESCRIPTION.md", content-type = "text/markdown" }
 description = "Protein design"
 dependencies = [
     # Add runtime dependencies here
-    "numpy==2.0.2",
-    "numba==0.61.0",
+    "numpy==2.0.2; platform_system != 'Darwin'",
+    "numba==0.61.0; platform_system != 'Darwin'",
     "matplotlib",
     "hydride",
     "biotite",
@@ -33,10 +33,10 @@ dependencies = [
     "einx",
     "einops",
     "mashumaro",
-    "nvidia-ml-py>=12.535.133",
-    "cuequivariance_ops_cu12>=0.5.0",
-    "cuequivariance_ops_torch_cu12>=0.5.0",
-    "cuequivariance_torch>=0.5.0",
+    "nvidia-ml-py>=12.535.133; platform_system != 'Darwin'",
+    "cuequivariance_ops_cu12>=0.5.0; platform_system != 'Darwin'",
+    "cuequivariance_ops_torch_cu12>=0.5.0; platform_system != 'Darwin'",
+    "cuequivariance_torch>=0.5.0; platform_system != 'Darwin'",
     "huggingface_hub",
     "biopython", 
 ]

src/boltzgen/cli/boltzgen.py

@@ -917,8 +917,8 @@ def __init__(self, args: argparse.Namespace, moldir: Path):
                 f"Invalid protocol: {protocol}. Valid protocols: {list(protocol_configs.keys())}"
             )
 
-        # Handle use_kernels argument
-        device_capability = torch.cuda.get_device_capability()
+        # Handle use_kernels argument, defaulting to (0,0) for CPU/MPS
+        device_capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else (0, 0)
         use_kernels = None
         if args.use_kernels == "auto":
             use_kernels = device_capability[0] >= 8
@@ -937,9 +937,9 @@ def __init__(self, args: argparse.Namespace, moldir: Path):
         config_args_by_step = parse_config_args(
             protocol_config, args.config, step_names
         )
-
+        # Determine number of devices to use, defaulting to 1 for MPS/CPU
         devices = (
-            args.devices if args.devices is not None else torch.cuda.device_count()
+            args.devices if args.devices is not None else torch.cuda.device_count() if torch.cuda.is_available() else 1
         )
         print(f"Using {devices} devices")
 

src/boltzgen/data/filter/static/polymer.py

@@ -142,6 +142,11 @@ def filter(self, structure: Structure) -> np.ndarray:
             res_start = chain["res_idx"]
             res_end = res_start + chain["res_num"]
             residues = structure.residues[res_start:res_end]
+            # Exclude non-standard residues (e.g. bound metal ions/hetero
+            # groups) so their atom_center isn't treated as a backbone CA,
+            # which would produce a spurious long jump and wrongly reject
+            # an otherwise valid chain.
+            residues = residues[residues["is_standard"]]
 
             # Get c-alphas
             ca_ids = residues["atom_center"]

src/boltzgen/model/layers/confidence_utils.py

@@ -22,7 +22,7 @@ def compute_frame_pred(
     resolved_mask=None,
     inference=False,
 ):
-    with torch.amp.autocast("cuda", enabled=False):
+    with torch.amp.autocast(device_type=pred_atom_coords.device.type, enabled=False):
         asym_id_token = feats["asym_id"]
         asym_id_atom = torch.bmm(
             feats["atom_to_token"].float(), asym_id_token.unsqueeze(-1).float()

src/boltzgen/model/loss/diffusion.py

@@ -80,9 +80,17 @@ def weighted_rigid_align(
     original_dtype = cov_matrix.dtype
     cov_matrix_32 = cov_matrix.to(dtype=torch.float32)
 
-    U, S, V = torch.linalg.svd(
-        cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None
-    )
+    # move cov_matrix_32 to cpu for mps compatibility
+    if cov_matrix_32.device.type == "mps":
+        cov_matrix_cpu = cov_matrix_32.cpu()
+        U, S, V = torch.linalg.svd(cov_matrix_cpu, driver=None)
+        U = U.to(cov_matrix_32.device)
+        S = S.to(cov_matrix_32.device)
+        V = V.to(cov_matrix_32.device)
+    else:
+        U, S, V = torch.linalg.svd(
+            cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None
+        )
     V = V.mH
 
     # Catch ambiguous rotation by checking the magnitude of singular values

src/boltzgen/model/models/boltz.py

@@ -655,7 +655,7 @@ def forward(
             ):
                 if self.inference_logging:
                     print("\nRunning Structure Module.\n")
-                with torch.autocast("cuda", enabled=False):
+                with torch.autocast(device_type=s.device.type, enabled=False):
                     if not self.inverse_fold:
                         struct_out = self.structure_module.sample(
                             s_trunk=s.float(),
@@ -711,7 +711,7 @@ def forward(
                 feats["coords"] = atom_coords  # (multiplicity, L, 3)
                 assert len(feats["coords"].shape) == 3
 
-                with torch.autocast("cuda", enabled=False):
+                with torch.autocast(device_type=atom_coords.device.type, enabled=False):
                     if not self.inverse_fold:
                         struct_out = self.structure_module(
                             s_trunk=s.float(),
@@ -769,7 +769,7 @@ def forward(
             ]
             s_inputs = self.input_embedder(feats, affinity=True)
 
-            with torch.autocast("cuda", enabled=False):
+            with torch.autocast(device_type=s_inputs.device.type, enabled=False):
                 if self.affinity_ensemble:
                     dict_out_affinity1 = self.affinity_module1(
                         s_inputs=s_inputs.detach(),
@@ -1112,7 +1112,7 @@ def parameter_norm(self, module):
         parameters = [p.norm(p=2) ** 2 for p in module.parameters() if p.requires_grad]
         if len(parameters) == 0:
             return torch.tensor(
-                0.0, device="cuda" if torch.cuda.is_available() else "cpu"
+                0.0, device="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
             )
         norm = torch.stack(parameters).sum().sqrt()
         return norm
@@ -1165,7 +1165,10 @@ def validation_step(
                         "res_type =",
                         batch["res_type"].shape,
                     )
-                    torch.cuda.empty_cache()
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                    if torch.backends.mps.is_available():
+                        torch.mps.empty_cache()
                     return
                 raise e
         else:
@@ -1184,7 +1187,10 @@ def validation_step(
                 if "out of memory" in str(e):
                     msg = f"| WARNING: ran out of memory, skipping batch, {idx_dataset}"
                     print(msg)
-                    torch.cuda.empty_cache()
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                    if torch.backends.mps.is_available():
+                        torch.mps.empty_cache()
                     return
                 raise e
 
@@ -1371,7 +1377,10 @@ def predict_step(
         except RuntimeError as e:  # catch out of memory exceptions
             if "out of memory" in str(e):
                 print("| WARNING: ran out of memory, skipping batch")
-                torch.cuda.empty_cache()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                if torch.backends.mps.is_available():
+                    torch.mps.empty_cache()
                 return {"exception": True}
             else:
                 raise e

src/boltzgen/model/modules/inverse_fold.py

@@ -470,7 +470,7 @@ def extract_attr_feat(
             dim=-1,
         )
 
-        with torch.autocast("cuda", enabled=False):
+        with torch.autocast(device_type=feats["atom_to_token"].device.type, enabled=False):
             atom_to_token = feats["atom_to_token"].float()
             atom_to_token_mean = atom_to_token / (
                 atom_to_token.sum(dim=1, keepdim=True) + 1e-6

src/boltzgen/model/modules/trunk.py

@@ -357,7 +357,7 @@ def forward(
         ).float()
 
         # Compute template features
-        with torch.autocast(device_type="cuda", enabled=False):
+        with torch.autocast(device_type=cb_coords.device.type, enabled=False):
             # Compute distogram
             cb_dists = torch.cdist(cb_coords, cb_coords)
             boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1)
@@ -507,7 +507,7 @@ def forward(
         token_coords = feats["center_coords"]
 
         # Compute template features
-        with torch.autocast(device_type="cuda", enabled=False):
+        with torch.autocast(device_type=token_coords.device.type, enabled=False):
             # Compute distogram
             dists = torch.cdist(token_coords, token_coords)
             boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1)

src/boltzgen/model/validation/refolding.py

@@ -299,11 +299,14 @@ def on_epoch_end(self, model):
         self.folding_model = None
         del self.affinity_model
         self.affinity_model = None
-        torch._C._cuda_clearCublasWorkspaces()
+        if torch.cuda.is_available():
+            torch._C._cuda_clearCublasWorkspaces()
         torch._dynamo.reset()
         gc.collect()
-        torch.cuda.empty_cache()
-
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        elif torch.backends.mps.is_available():
+            torch.mps.empty_cache()
         # Compute standard metrics
         self.common_on_epoch_end(model, logname="val_monomer_ligand")
         self.on_epoch_end_design(model, logname="val_monomer_ligand")

src/boltzgen/task/predict/data_from_generated.py

@@ -427,16 +427,17 @@ def get_feat(self, path, design_mask, ss_type=None, binding_type=None, aa_constr
         try:
             # Try to find molecules in the dataset moldir if provided
             # Find missing ones in global moldir and check if all found
+            # Note: load fresh from moldir to avoid losing RDKit atom properties
+            # when self.canonicals is pickled by DataLoader worker processes.
             molecules = {}
-            molecules.update(self.canonicals)
             mol_names = set(tokenized.tokens["res_name"].tolist())
-            mol_names = mol_names - set(self.canonicals.keys())
             if mols is not None:
                 molecules.update(mols)
-            mol_names = mol_names - set(molecules.keys())
+                mol_names = mol_names - set(mols.keys())
             if self.moldir is not None:
                 molecules.update(load_molecules(self.moldir, mol_names))
-            molecules.update(load_molecules(self.moldir, mol_names))
+            else:
+                molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names})
         except Exception as e:  # noqa: BLE001
             print(f"Molecule loading failed for {path} with error {e}. Skipping.")
             raise DataFetchException() from e
@@ -840,11 +841,13 @@ def output_path_analyzed(input_path):
         )
 
     def predict_dataloader(self) -> DataLoader:
+        pin_memory = self.cfg.pin_memory and not torch.backends.mps.is_available()
         return DataLoader(
             self.predict_set,
             batch_size=self.cfg.batch_size,
             num_workers=self.cfg.num_workers,
-            pin_memory=self.cfg.pin_memory,
+            pin_memory=pin_memory,
+            persistent_workers=self.cfg.num_workers > 0,
             shuffle=False,
             collate_fn=collate,
         )
@@ -877,6 +880,8 @@ def transfer_batch_to_device(
                 "tokenized",
                 "data_sample_idx",
             ]:
+                if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available():
+                    batch[key] = batch[key].float()
                 batch[key] = batch[key].to(device)
 
         return batch

src/boltzgen/task/predict/data_from_yaml.py

@@ -237,16 +237,16 @@ def get_sample(self, path: Path, sample_id: Optional[str] = None) -> Dict:
 
         # Try to find molecules in the dataset moldir if provided
         # Find missing ones in global moldir and check if all found
+        # Note: self.canonicals may lose atom-level RDKit properties when pickled
+        # by the DataLoader for worker processes (Python pickle does not preserve
+        # RDKit atom SetProp values). Load all molecules fresh from moldir instead.
         molecules = {}
-        molecules.update(self.canonicals)
         mol_names = set(tokenized.tokens["res_name"].tolist())
-        mol_names = mol_names - set(self.canonicals.keys())
         mol_names = mol_names - set(parsed.extra_mols.keys())
         if self.moldir is not None:
             molecules.update(load_molecules(self.moldir, mol_names))
-
-        mol_names = mol_names - set(molecules.keys())
-        molecules.update(load_molecules(self.moldir, mol_names))
+        else:
+            molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names})
         molecules.update(parsed.extra_mols)
 
         # Finalize input data
@@ -393,11 +393,13 @@ def predict_dataloader(self) -> DataLoader:
             The training dataloader.
 
         """
+        pin_memory = self.pin_memory and not torch.backends.mps.is_available()
         return DataLoader(
             self.predict_set,
             batch_size=self.batch_size,
             num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
+            pin_memory=pin_memory,
+            persistent_workers=self.num_workers > 0,
             shuffle=False,
             collate_fn=collate,
         )
@@ -445,5 +447,8 @@ def transfer_batch_to_device(
                 "extra_mols",
                 "data_sample_idx",
             ]:
+                #Convert torch.float64 to torch.float for mps compatibility
+                if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available():
+                    batch[key] = batch[key].float()
                 batch[key] = batch[key].to(device)
         return batch

src/boltzgen/task/predict/data_ligands.py

@@ -221,15 +221,14 @@ def __getitem__(self, idx: int) -> Dict:
         try:
             # Try to find molecules in the dataset moldir if provided
             # Find missing ones in global moldir and check if all found
+            # Note: load fresh from moldir to avoid losing RDKit atom properties
+            # when self.canonicals is pickled by DataLoader worker processes.
             molecules = {}
-            molecules.update(self.canonicals)
             mol_names = set(tokenized.tokens["res_name"].tolist())
-            mol_names = mol_names - set(self.canonicals.keys())
             if self.moldir is not None:
                 molecules.update(load_molecules(self.moldir, mol_names))
-
-            mol_names = mol_names - set(molecules.keys())
-            molecules.update(load_molecules(self.moldir, mol_names))
+            else:
+                molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names})
         except Exception as e:  # noqa: BLE001
             print(f"Molecule loading failed for {target_id} with error {e}. Skipping.")
             return self.__getitem__(0)
@@ -352,11 +351,13 @@ def predict_dataloader(self) -> DataLoader:
             The training dataloader.
 
         """
+        pin_memory = self.pin_memory and not torch.backends.mps.is_available()
         return DataLoader(
             self.predict_set,
             batch_size=self.batch_size,
             num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
+            pin_memory=pin_memory,
+            persistent_workers=self.num_workers > 0,
             shuffle=False,
             collate_fn=collate,
         )
@@ -406,5 +407,7 @@ def transfer_batch_to_device(
                 "extra_mols",
                 "data_sample_idx",
             ]:
+                if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available():
+                    batch[key] = batch[key].float()
                 batch[key] = batch[key].to(device)
         return batch