boltz-vsynthes 0.0.13__tar.gz → 0.0.14__tar.gz

{boltz_vsynthes-0.0.13 → boltz_vsynthes-0.0.14}/PKG-INFO

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

{boltz_vsynthes-0.0.13 → boltz_vsynthes-0.0.14}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "boltz-vsynthes"
-version = "0.0.13"
+version = "0.0.14"
 requires-python = ">=3.10,<3.13"
 description = "Boltz for VSYNTHES"
 readme = "README.md"

{boltz_vsynthes-0.0.13 → boltz_vsynthes-0.0.14}/src/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 → boltz_vsynthes-0.0.14}/src/boltz_vsynthes.egg-info/PKG-INFO

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

{boltz_vsynthes-0.0.13 → boltz_vsynthes-0.0.14}/src/boltz_vsynthes.egg-info/SOURCES.txt

@@ -1,7 +1,6 @@
 LICENSE
 README.md
 pyproject.toml
-setup.py
 src/boltz/__init__.py
 src/boltz/main.py
 src/boltz/data/__init__.py

boltz_vsynthes-0.0.13/setup.py

@@ -1,10 +0,0 @@
-from setuptools import setup, find_packages
-
-setup(
-    # ... existing setup args ...
-    entry_points={
-        'console_scripts': [
-            'boltz-sdf-to-pre-affinity=boltz.utils.sdf_to_pre_affinity_npz:main',
-        ],
-    },
-)