boltz-vsynthes 0.0.17__py3-none-any.whl → 0.1.0__py3-none-any.whl

boltz/data/module/inferencev2.py

@@ -59,16 +59,16 @@ def load_input(
 
     """
     # Load the structure
-    # if affinity:
-    #     structure = StructureV2.load(
-    #         target_dir / record.id / f"pre_affinity_{record.id}.npz"
-    #     )
     if affinity:
-        if target_dir.name == "predictions":
-            target_dir = target_dir.parent / "processed"
         structure = StructureV2.load(
-            target_dir / f"structures/{record.id}.npz"
+            target_dir / record.id / f"pre_affinity_{record.id}.npz"
         )
+    # if affinity:
+    #     if target_dir.name == "predictions":
+    #         target_dir = target_dir.parent / "processed"
+    #     structure = StructureV2.load(
+    #         target_dir / f"structures/{record.id}.npz"
+    #     )
     else:
         structure = StructureV2.load(target_dir / f"{record.id}.npz")
 

boltz/main.py

@@ -9,7 +9,7 @@ from dataclasses import asdict, dataclass
 from functools import partial
 from multiprocessing import Pool
 from pathlib import Path
-from typing import Literal, Optional
+from typing import Literal, Optional, List
 
 import click
 import torch
@@ -18,6 +18,8 @@ from pytorch_lightning.strategies import DDPStrategy
 from pytorch_lightning.utilities import rank_zero_only
 from rdkit import Chem
 from tqdm import tqdm
+import time
+from datetime import datetime
 
 from boltz.data import const
 from boltz.data.module.inference import BoltzInferenceDataModule
@@ -203,22 +205,21 @@ def download_boltz2(cache: Path) -> None:
         The cache directory.
 
     """
+    # Use /tmp if possible for faster local disk I/O
+    if str(cache).startswith("/home") or str(cache).startswith("/mnt"):
+        cache = Path("/tmp/boltz_cache")
+    cache.mkdir(parents=True, exist_ok=True)
+
     # Download CCD
     mols = cache / "mols"
     tar_mols = cache / "mols.tar"
-    if not tar_mols.exists():
-        click.echo(
-            f"Downloading the CCD data to {tar_mols}. "
-            "This may take a bit of time. You may change the cache directory "
-            "with the --cache flag."
-        )
-        urllib.request.urlretrieve(MOL_URL, str(tar_mols))  # noqa: S310
     if not mols.exists():
         click.echo(
-            f"Extracting the CCD data to {mols}. "
+            f"Downloading and extracting the CCD data to {mols}. "
             "This may take a bit of time. You may change the cache directory "
             "with the --cache flag."
         )
+        urllib.request.urlretrieve(MOL_URL, str(tar_mols))  # noqa: S310
         with tarfile.open(str(tar_mols), "r") as tar:
             tar.extractall(cache)  # noqa: S202
 
@@ -983,6 +984,7 @@ def predict(  # noqa: C901, PLR0915, PLR0912
     torch.set_grad_enabled(False)
 
     # Ignore matmul precision warning
+    # torch.set_float32_matmul_precision('medium')
     torch.set_float32_matmul_precision("highest")
 
     # Set rdkit pickle logic
@@ -1029,14 +1031,16 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             msg = f"Method {method} not supported. Supported: {method_names}"
             raise ValueError(msg)
 
-    # Process inputs
+    # 1. Before and after process_inputs
+    t_process_inputs = time.time()
     ccd_path = cache / "ccd.pkl"
     mol_dir = cache / "mols"
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Starting process_inputs...")
     process_inputs(
         data=data,
         out_dir=out_dir,
-        ccd_path=ccd_path,
-        mol_dir=mol_dir,
+        ccd_path=cache / "ccd.pkl",
+        mol_dir=cache / "mols",
         use_msa_server=use_msa_server,
         msa_server_url=msa_server_url,
         msa_pairing_strategy=msa_pairing_strategy,
@@ -1044,18 +1048,27 @@ def predict(  # noqa: C901, PLR0915, PLR0912
         preprocessing_threads=preprocessing_threads,
         max_msa_seqs=max_msa_seqs,
     )
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] process_inputs finished in {time.time() - t_process_inputs:.2f} seconds")
 
-    # Load manifest
+    # 2. Before and after load manifest
+    t_manifest = time.time()
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Loading manifest...")
     manifest = Manifest.load(out_dir / "processed" / "manifest.json")
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Manifest loaded in {time.time() - t_manifest:.2f} seconds")
 
-    # Filter out existing predictions
+    # 3. Before and after Filter out existing predictions
+    t_filter = time.time()
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Filtering out existing predictions...")
     filtered_manifest = filter_inputs_structure(
         manifest=manifest,
         outdir=out_dir,
         override=override,
     )
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Filtering finished in {time.time() - t_filter:.2f} seconds")
 
-    # Load processed data
+    # 4. Before and after load processed data
+    t_processed = time.time()
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Loading processed data...")
     processed_dir = out_dir / "processed"
     processed = BoltzProcessedInput(
         manifest=filtered_manifest,
@@ -1075,8 +1088,22 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             (processed_dir / "mols") if (processed_dir / "mols").exists() else None
         ),
     )
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Processed data loaded in {time.time() - t_processed:.2f} seconds")
 
-    # Set up trainer
+    # 5. Before and after create prediction writer
+    t_writer = time.time()
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Creating prediction writer...")
+    pred_writer = BoltzWriter(
+        data_dir=processed.targets_dir,
+        output_dir=out_dir / "predictions",
+        output_format=output_format,
+        boltz2=model == "boltz2",
+    )
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Prediction writer created in {time.time() - t_writer:.2f} seconds")
+
+    # 6. Before and after set up trainer
+    t_trainer = time.time()
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Setting up trainer...")
     strategy = "auto"
     if (isinstance(devices, int) and devices > 1) or (
         isinstance(devices, list) and len(devices) > 1
@@ -1129,13 +1156,16 @@ def predict(  # noqa: C901, PLR0915, PLR0912
         devices=devices,
         precision=32 if model == "boltz1" else "bf16-mixed",
     )
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] Trainer set up in {time.time() - t_trainer:.2f} seconds")
 
     if filtered_manifest.records:
         msg = f"Running structure prediction for {len(filtered_manifest.records)} input"
         msg += "s." if len(filtered_manifest.records) > 1 else "."
         click.echo(msg)
 
-        # Create data module
+        # 7. Before and after create data module
+        t_datamodule = time.time()
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Creating data module...")
         if model == "boltz2":
             data_module = Boltz2InferenceDataModule(
                 manifest=processed.manifest,
@@ -1156,14 +1186,16 @@ def predict(  # noqa: C901, PLR0915, PLR0912
                 num_workers=num_workers,
                 constraints_dir=processed.constraints_dir,
             )
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Data module created in {time.time() - t_datamodule:.2f} seconds")
 
-        # Load model
+        # 8. Before and after load model
+        t_model = time.time()
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Loading model...")
         if checkpoint is None:
             if model == "boltz2":
                 checkpoint = cache / "boltz2_conf.ckpt"
             else:
                 checkpoint = cache / "boltz1_conf.ckpt"
-
         predict_args = {
             "recycling_steps": recycling_steps,
             "sampling_steps": sampling_steps,
@@ -1173,12 +1205,10 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             "write_full_pae": write_full_pae,
             "write_full_pde": write_full_pde,
         }
-
         steering_args = BoltzSteeringParams()
         steering_args.fk_steering = use_potentials
         steering_args.guidance_update = use_potentials
-
-        model_cls = Boltz2 if model == "boltz2" else Boltz1
+        model_cls = Boltz2 if model == "boltz2" else Boltz1)
         model_module = model_cls.load_from_checkpoint(
             checkpoint,
             strict=True,
@@ -1192,13 +1222,17 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             steering_args=asdict(steering_args),
         )
         model_module.eval()
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Model loaded in {time.time() - t_model:.2f} seconds")
 
-        # Compute structure predictions
+        # 9. Before and after compute structure predictions (predict)
+        t_predict = time.time()
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Computing structure predictions...")
         trainer.predict(
             model_module,
             datamodule=data_module,
             return_predictions=False,
         )
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Structure predictions computed in {time.time() - t_predict:.2f} seconds")
 
     # Check if affinity predictions are needed
     if any(r.affinity for r in manifest.records):
@@ -1251,6 +1285,8 @@ def predict(  # noqa: C901, PLR0915, PLR0912
         if affinity_checkpoint is None:
             affinity_checkpoint = cache / "boltz2_aff.ckpt"
 
+        # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"[{datetime.now().strftime('%H:%M:%S')}] Using device: {device}")
         model_module = Boltz2.load_from_checkpoint(
             affinity_checkpoint,
             strict=True,
@@ -1272,6 +1308,5 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             return_predictions=False,
         )
 
-
 if __name__ == "__main__":
     cli()

boltz/model/models/boltz2.py

@@ -63,7 +63,7 @@ class Boltz2(LightningModule):
         num_val_datasets: int = 1,
         atom_feature_dim: int = 128,
         template_args: Optional[dict] = None,
-        confidence_prediction: bool = True,
+        confidence_prediction: bool = True, #TODO: change to False
         affinity_prediction: bool = False,
         affinity_ensemble: bool = False,
         affinity_mw_correction: bool = True,

{boltz_vsynthes-0.0.17.dist-info → boltz_vsynthes-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: boltz-vsynthes
-Version: 0.0.17
+Version: 0.1.0
 Summary: Boltz for VSYNTHES
 Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
@@ -28,9 +28,9 @@ Requires-Dist: numba==0.61.0
 Requires-Dist: gemmi==0.6.5
 Requires-Dist: scikit-learn==1.6.1
 Requires-Dist: chembl_structure_pipeline==1.2.2
-Requires-Dist: cuequivariance_ops_cu12>=0.5.0
-Requires-Dist: cuequivariance_ops_torch_cu12>=0.5.0
-Requires-Dist: cuequivariance_torch>=0.5.0
+Requires-Dist: cuequivariance_ops_cu12==0.5.0
+Requires-Dist: cuequivariance_ops_torch_cu12==0.5.0
+Requires-Dist: cuequivariance_torch==0.5.0
 Provides-Extra: lint
 Requires-Dist: ruff; extra == "lint"
 Provides-Extra: test

{boltz_vsynthes-0.0.17.dist-info → boltz_vsynthes-0.1.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
 boltz/__init__.py,sha256=F_-so3S40iZrSZ89Ge4TS6aZqwWyZXq_H4AXGDlbA_g,187
-boltz/main.py,sha256=zBLxa6T8hxcBs7gj1BnWfgJSx6uki8iV-QgClvoaiSA,39951
+boltz/main.py,sha256=BApAG6y3m_V5RuwTtoBL_f2Ud69BVPT3pPSfsG1R718,42706
 boltz/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/data/const.py,sha256=1M-88Z6HkfKY6MkNtqcj3b9P-oX9xEXluh3qM_u8dNU,26779
 boltz/data/mol.py,sha256=maOpPHEGX1VVXCIFY6pQNGF7gUBZPAfgSvuPf2QO1yc,34268
@@ -27,7 +27,7 @@ boltz/data/filter/static/ligand.py,sha256=LamC-Z9IjYj3DmfxwMFmPbKBBhRMby3uWQj74w
 boltz/data/filter/static/polymer.py,sha256=LNsQMsOOnhYpeKuM9AStktoTQPMZE3H0yu4mRg-jwPc,9386
 boltz/data/module/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/data/module/inference.py,sha256=xk8ZJ8UhjPiPTdOluH_v4gnV8GtTX3sr1WZ1s5Ox8I8,8100
-boltz/data/module/inferencev2.py,sha256=QmtvyUuHkY2f1ulE2xkIjo87xRIuOj3Yddo6_mbILYg,12738
+boltz/data/module/inferencev2.py,sha256=aLUm1WR6E1814JUrF6sJfoe5y8y7d_s4zlQ3pdFBVy8,12742
 boltz/data/module/training.py,sha256=iNzmq9ufs20S4M947CCzdYzGTFjmCTf2tFExJ2PtXnA,22428
 boltz/data/module/trainingv2.py,sha256=ZsYUHYXxfuPgIpbTwCj5QLO0XK__xjsqIw6GARSNGW0,21276
 boltz/data/msa/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -85,7 +85,7 @@ boltz/model/loss/distogramv2.py,sha256=dFgMGwpdLK4-skHJwvpERG10KfF3ZUN1T9_hUj-iW
 boltz/model/loss/validation.py,sha256=gYpbag9mulg5HJPXjOUFaMV9XSYX_s2bIQ0iYjiAow0,33501
 boltz/model/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/model/models/boltz1.py,sha256=x-x0b3VAXiAkPUBePnF56k1aYEPNgX1M6GtNCYVdCso,51718
-boltz/model/models/boltz2.py,sha256=3XOWjUWaSJquw8Xdp7JItDUnVDyoC0qtx3q4MFQrd38,51523
+boltz/model/models/boltz2.py,sha256=hD1kF4XSox9PPeLN_v02YrMg8zFnGyI8ZYW_AwOGJnc,51546
 boltz/model/modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/model/modules/affinity.py,sha256=FktI2wrkDqsjGHJOuvzVrZK78MOPjU65QN0l6sB1QPQ,7041
 boltz/model/modules/confidence.py,sha256=sXGymZiiMtfXPkUvHpa2KCCvNY79D8jXXEx9Gz2rNFs,17475
@@ -110,9 +110,9 @@ boltz/model/potentials/schedules.py,sha256=m7XJjfuF9uTX3bR9VisXv1rvzJjxiD8PobXRp
 boltz/utils/sdf_splitter.py,sha256=ZHn_syOcmm-fDnJ3YEGyGv_vYz2IRzUW7vbbMSU2JBY,2108
 boltz/utils/sdf_to_pre_affinity_npz.py,sha256=ro0KGe24JexbJm47J8S8w8Lmr_KaQbzOAb_dKZO2G9I,40384
 boltz/utils/yaml_generator.py,sha256=ermWIG-BE6nNWHFvpEwpk92N9J-YATpGXZGLvD1I2oQ,4012
-boltz_vsynthes-0.0.17.dist-info/licenses/LICENSE,sha256=8GZ_1eZsUeG6jdqgJJxtciWzADfgLEV4LY8sKUOsJhc,1102
-boltz_vsynthes-0.0.17.dist-info/METADATA,sha256=C3WBB1KJrsEeeiuDnk78uhHAdFRMp8ApLNnOS8uLx1c,7235
-boltz_vsynthes-0.0.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-boltz_vsynthes-0.0.17.dist-info/entry_points.txt,sha256=nZNYPKKrmAr-MVA0K-ClNRT2p90FV1_14d7HpsESZFQ,211
-boltz_vsynthes-0.0.17.dist-info/top_level.txt,sha256=MgU3Jfb-ctWm07YGMts68PMjSh9v26D0gfG3dFRmVFA,6
-boltz_vsynthes-0.0.17.dist-info/RECORD,,
+boltz_vsynthes-0.1.0.dist-info/licenses/LICENSE,sha256=8GZ_1eZsUeG6jdqgJJxtciWzADfgLEV4LY8sKUOsJhc,1102
+boltz_vsynthes-0.1.0.dist-info/METADATA,sha256=KKvZxnnHx3XX6WsLY4WA8lE2bY4OaTZa4RnYSP_rsfs,7234
+boltz_vsynthes-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+boltz_vsynthes-0.1.0.dist-info/entry_points.txt,sha256=nZNYPKKrmAr-MVA0K-ClNRT2p90FV1_14d7HpsESZFQ,211
+boltz_vsynthes-0.1.0.dist-info/top_level.txt,sha256=MgU3Jfb-ctWm07YGMts68PMjSh9v26D0gfG3dFRmVFA,6
+boltz_vsynthes-0.1.0.dist-info/RECORD,,