boltz-vsynthes 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl

boltz/data/module/inferencev2.py

@@ -59,9 +59,15 @@ 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 / record.id / f"pre_affinity_{record.id}.npz"
+            target_dir / record.id / f"{record.id}.npz"
         )
     else:
         structure = StructureV2.load(target_dir / f"{record.id}.npz")

boltz/utils/sdf_to_pre_affinity_npz.py

@@ -1076,41 +1076,41 @@ def predict(  # noqa: C901, PLR0915, PLR0912
         ),
     )
 
-    # # Set up trainer
-    # strategy = "auto"
-    # if (isinstance(devices, int) and devices > 1) or (
-    #     isinstance(devices, list) and len(devices) > 1
-    # ):
-    #     start_method = "fork" if platform.system() != "win32" else "spawn"
-    #     strategy = DDPStrategy(start_method=start_method)
-    #     if len(filtered_manifest.records) < devices:
-    #         msg = (
-    #             "Number of requested devices is greater "
-    #             "than the number of predictions, taking the minimum."
-    #         )
-    #         click.echo(msg)
-    #         if isinstance(devices, list):
-    #             devices = devices[: max(1, len(filtered_manifest.records))]
-    #         else:
-    #             devices = max(1, min(len(filtered_manifest.records), devices))
-
-    # # Set up model parameters
-    # if model == "boltz2":
-    #     diffusion_params = Boltz2DiffusionParams()
-    #     step_scale = 1.5 if step_scale is None else step_scale
-    #     diffusion_params.step_scale = step_scale
-    #     pairformer_args = PairformerArgsV2()
-    # else:
-    #     diffusion_params = BoltzDiffusionParams()
-    #     step_scale = 1.638 if step_scale is None else step_scale
-    #     diffusion_params.step_scale = step_scale
-    #     pairformer_args = PairformerArgs()
-
-    # msa_args = MSAModuleArgs(
-    #     subsample_msa=subsample_msa,
-    #     num_subsampled_msa=num_subsampled_msa,
-    #     use_paired_feature=model == "boltz2",
-    # )
+    # Set up trainer
+    strategy = "auto"
+    if (isinstance(devices, int) and devices > 1) or (
+        isinstance(devices, list) and len(devices) > 1
+    ):
+        start_method = "fork" if platform.system() != "win32" else "spawn"
+        strategy = DDPStrategy(start_method=start_method)
+        if len(filtered_manifest.records) < devices:
+            msg = (
+                "Number of requested devices is greater "
+                "than the number of predictions, taking the minimum."
+            )
+            click.echo(msg)
+            if isinstance(devices, list):
+                devices = devices[: max(1, len(filtered_manifest.records))]
+            else:
+                devices = max(1, min(len(filtered_manifest.records), devices))
+
+    # Set up model parameters
+    if model == "boltz2":
+        diffusion_params = Boltz2DiffusionParams()
+        step_scale = 1.5 if step_scale is None else step_scale
+        diffusion_params.step_scale = step_scale
+        pairformer_args = PairformerArgsV2()
+    else:
+        diffusion_params = BoltzDiffusionParams()
+        step_scale = 1.638 if step_scale is None else step_scale
+        diffusion_params.step_scale = step_scale
+        pairformer_args = PairformerArgs()
+
+    msa_args = MSAModuleArgs(
+        subsample_msa=subsample_msa,
+        num_subsampled_msa=num_subsampled_msa,
+        use_paired_feature=model == "boltz2",
+    )
 
     # # Create prediction writer
     # pred_writer = BoltzWriter(
@@ -1200,77 +1200,86 @@ def predict(  # noqa: C901, PLR0915, PLR0912
     #         return_predictions=False,
     #     )
 
-    # # Check if affinity predictions are needed
-    # if any(r.affinity for r in manifest.records):
-    #     # Print header
-    #     click.echo("\nPredicting property: affinity\n")
+    # Check if affinity predictions are needed
+    if any(r.affinity for r in manifest.records):
+        # Print header
+        click.echo("\nPredicting property: affinity\n")
 
-    #     # Validate inputs
-    #     manifest_filtered = filter_inputs_affinity(
-    #         manifest=manifest,
-    #         outdir=out_dir,
-    #         override=override,
-    #     )
-    #     if not manifest_filtered.records:
-    #         click.echo("Found existing affinity predictions for all inputs, skipping.")
-    #         return
-
-    #     msg = f"Running affinity prediction for {len(manifest_filtered.records)} input"
-    #     msg += "s." if len(manifest_filtered.records) > 1 else "."
-    #     click.echo(msg)
+        # Validate inputs
+        manifest_filtered = filter_inputs_affinity(
+            manifest=manifest,
+            outdir=out_dir,
+            override=override,
+        )
+        if not manifest_filtered.records:
+            click.echo("Found existing affinity predictions for all inputs, skipping.")
+            return
 
-    #     pred_writer = BoltzAffinityWriter(
-    #         data_dir=processed.targets_dir,
-    #         output_dir=out_dir / "predictions",
-    #     )
+        msg = f"Running affinity prediction for {len(manifest_filtered.records)} input"
+        msg += "s." if len(manifest_filtered.records) > 1 else "."
+        click.echo(msg)
 
-    #     data_module = Boltz2InferenceDataModule(
-    #         manifest=manifest_filtered,
-    #         target_dir=out_dir / "predictions",
-    #         msa_dir=processed.msa_dir,
-    #         mol_dir=mol_dir,
-    #         num_workers=num_workers,
-    #         constraints_dir=processed.constraints_dir,
-    #         template_dir=processed.template_dir,
-    #         extra_mols_dir=processed.extra_mols_dir,
-    #         override_method="other",
-    #         affinity=True,
-    #     )
+        pred_writer = BoltzAffinityWriter(
+            data_dir=processed.targets_dir,
+            output_dir=out_dir / "predictions",
+        )
 
-    #     predict_affinity_args = {
-    #         "recycling_steps": 5,
-    #         "sampling_steps": sampling_steps_affinity,
-    #         "diffusion_samples": diffusion_samples_affinity,
-    #         "max_parallel_samples": 1,
-    #         "write_confidence_summary": False,
-    #         "write_full_pae": False,
-    #         "write_full_pde": False,
-    #     }
+        trainer = Trainer(
+            default_root_dir=out_dir,
+            strategy=strategy,
+            callbacks=[pred_writer],
+            accelerator=accelerator,
+            devices=devices,
+            precision=32 if model == "boltz1" else "bf16-mixed",
+        )
 
-    #     # Load affinity model
-    #     if affinity_checkpoint is None:
-    #         affinity_checkpoint = cache / "boltz2_aff.ckpt"
+        data_module = Boltz2InferenceDataModule(
+            manifest=manifest_filtered,
+            target_dir=out_dir / "predictions",
+            msa_dir=processed.msa_dir,
+            mol_dir=mol_dir,
+            num_workers=num_workers,
+            constraints_dir=processed.constraints_dir,
+            template_dir=processed.template_dir,
+            extra_mols_dir=processed.extra_mols_dir,
+            override_method="other",
+            affinity=True,
+        )
 
-    #     model_module = Boltz2.load_from_checkpoint(
-    #         affinity_checkpoint,
-    #         strict=True,
-    #         predict_args=predict_affinity_args,
-    #         map_location="cpu",
-    #         diffusion_process_args=asdict(diffusion_params),
-    #         ema=False,
-    #         pairformer_args=asdict(pairformer_args),
-    #         msa_args=asdict(msa_args),
-    #         steering_args={"fk_steering": False, "guidance_update": False},
-    #         affinity_mw_correction=affinity_mw_correction,
-    #     )
-    #     model_module.eval()
+        predict_affinity_args = {
+            "recycling_steps": 5,
+            "sampling_steps": sampling_steps_affinity,
+            "diffusion_samples": diffusion_samples_affinity,
+            "max_parallel_samples": 1,
+            "write_confidence_summary": False,
+            "write_full_pae": False,
+            "write_full_pde": False,
+        }
+
+        # Load affinity model
+        if affinity_checkpoint is None:
+            affinity_checkpoint = cache / "boltz2_aff.ckpt"
+
+        model_module = Boltz2.load_from_checkpoint(
+            affinity_checkpoint,
+            strict=True,
+            predict_args=predict_affinity_args,
+            map_location="cpu",
+            diffusion_process_args=asdict(diffusion_params),
+            ema=False,
+            pairformer_args=asdict(pairformer_args),
+            msa_args=asdict(msa_args),
+            steering_args={"fk_steering": False, "guidance_update": False},
+            affinity_mw_correction=affinity_mw_correction,
+        )
+        model_module.eval()
 
-    #     trainer.callbacks[0] = pred_writer
-    #     trainer.predict(
-    #         model_module,
-    #         datamodule=data_module,
-    #         return_predictions=False,
-    #     )
+        trainer.callbacks[0] = pred_writer
+        trainer.predict(
+            model_module,
+            datamodule=data_module,
+            return_predictions=False,
+        )
 
 
 if __name__ == "__main__":

{boltz_vsynthes-0.0.13.dist-info → boltz_vsynthes-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: boltz-vsynthes
-Version: 0.0.13
+Version: 0.0.15
 Summary: Boltz for VSYNTHES
 Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown

{boltz_vsynthes-0.0.13.dist-info → boltz_vsynthes-0.0.15.dist-info}/RECORD

@@ -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=3p-jyPstcNzUeaOshEzAexXRBFjvpr-9tP3n8hxT6nw,12508
+boltz/data/module/inferencev2.py,sha256=y1bN4EchiQ3BsWd6Tz6Ug41I_TItP6WQTc4amBuRN0E,12739
 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
@@ -108,11 +108,11 @@ boltz/model/potentials/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3
 boltz/model/potentials/potentials.py,sha256=vev8Vjfs-ML1hyrdv_R8DynG4wSFahJ6nzPWp7CYQqw,17507
 boltz/model/potentials/schedules.py,sha256=m7XJjfuF9uTX3bR9VisXv1rvzJjxiD8PobXRpcBBu1c,968
 boltz/utils/sdf_splitter.py,sha256=ZHn_syOcmm-fDnJ3YEGyGv_vYz2IRzUW7vbbMSU2JBY,2108
-boltz/utils/sdf_to_pre_affinity_npz.py,sha256=ENljAVhA7ZtDkUCp1xuJufTyVbuaQHZAe_vAl6ck-WE,40301
+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.13.dist-info/licenses/LICENSE,sha256=8GZ_1eZsUeG6jdqgJJxtciWzADfgLEV4LY8sKUOsJhc,1102
-boltz_vsynthes-0.0.13.dist-info/METADATA,sha256=-ZCZOOLVwXKOYf2E0XY8Q1ZAlHQgZFp0Qj4_mPfeIqU,7235
-boltz_vsynthes-0.0.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-boltz_vsynthes-0.0.13.dist-info/entry_points.txt,sha256=nZNYPKKrmAr-MVA0K-ClNRT2p90FV1_14d7HpsESZFQ,211
-boltz_vsynthes-0.0.13.dist-info/top_level.txt,sha256=MgU3Jfb-ctWm07YGMts68PMjSh9v26D0gfG3dFRmVFA,6
-boltz_vsynthes-0.0.13.dist-info/RECORD,,
+boltz_vsynthes-0.0.15.dist-info/licenses/LICENSE,sha256=8GZ_1eZsUeG6jdqgJJxtciWzADfgLEV4LY8sKUOsJhc,1102
+boltz_vsynthes-0.0.15.dist-info/METADATA,sha256=NJSqvHhTsV6-kN89Wa1LdFW8gmXutClh54r_iKsjo94,7235
+boltz_vsynthes-0.0.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+boltz_vsynthes-0.0.15.dist-info/entry_points.txt,sha256=nZNYPKKrmAr-MVA0K-ClNRT2p90FV1_14d7HpsESZFQ,211
+boltz_vsynthes-0.0.15.dist-info/top_level.txt,sha256=MgU3Jfb-ctWm07YGMts68PMjSh9v26D0gfG3dFRmVFA,6
+boltz_vsynthes-0.0.15.dist-info/RECORD,,