boltz-vsynthes 1.0.34__tar.gz → 1.0.35__tar.gz

{boltz_vsynthes-1.0.34/src/boltz_vsynthes.egg-info → boltz_vsynthes-1.0.35}/PKG-INFO

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

{boltz_vsynthes-1.0.34 → boltz_vsynthes-1.0.35}/pyproject.toml

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

{boltz_vsynthes-1.0.34 → boltz_vsynthes-1.0.35}/src/boltz/data/write/writer.py

@@ -290,35 +290,72 @@ class BoltzAffinityWriter(BasePredictionWriter):
         if prediction["exception"]:
             self.failed += 1
             return
-        # Dump affinity summary
-        affinity_summary = {}
-        pred_affinity_value = prediction["affinity_pred_value"]
-        pred_affinity_probability = prediction["affinity_probability_binary"]
-        affinity_summary = {
-            "affinity_pred_value": pred_affinity_value.item(),
-            "affinity_probability_binary": pred_affinity_probability.item(),
-        }
-        if "affinity_pred_value1" in prediction:
-            pred_affinity_value1 = prediction["affinity_pred_value1"]
-            pred_affinity_probability1 = prediction["affinity_probability_binary1"]
-            pred_affinity_value2 = prediction["affinity_pred_value2"]
-            pred_affinity_probability2 = prediction["affinity_probability_binary2"]
-            affinity_summary["affinity_pred_value1"] = pred_affinity_value1.item()
-            affinity_summary["affinity_probability_binary1"] = (
-                pred_affinity_probability1.item()
-            )
-            affinity_summary["affinity_pred_value2"] = pred_affinity_value2.item()
-            affinity_summary["affinity_probability_binary2"] = (
-                pred_affinity_probability2.item()
-            )
-
-        # Save the affinity summary
-        struct_dir = self.output_dir / batch["record"][0].id
-        struct_dir.mkdir(exist_ok=True)
-        path = struct_dir / f"affinity_{batch['record'][0].id}.json"
-
-        with path.open("w") as f:
-            f.write(json.dumps(affinity_summary, indent=4))
+
+        # Get the record and check if it has multiple ligands
+        record = batch["record"][0]
+        record_dir = self.output_dir / record.id
+
+        # Check if we have ligand-specific directories
+        ligand_dirs = [d for d in record_dir.iterdir() if d.is_dir() and d.name.startswith("LIG")]
+        
+        if ligand_dirs:
+            # We have multiple ligands, save predictions in each ligand's directory
+            for ligand_dir in ligand_dirs:
+                # Dump affinity summary
+                affinity_summary = {}
+                pred_affinity_value = prediction["affinity_pred_value"]
+                pred_affinity_probability = prediction["affinity_probability_binary"]
+                affinity_summary = {
+                    "affinity_pred_value": pred_affinity_value.item(),
+                    "affinity_probability_binary": pred_affinity_probability.item(),
+                }
+                if "affinity_pred_value1" in prediction:
+                    pred_affinity_value1 = prediction["affinity_pred_value1"]
+                    pred_affinity_probability1 = prediction["affinity_probability_binary1"]
+                    pred_affinity_value2 = prediction["affinity_pred_value2"]
+                    pred_affinity_probability2 = prediction["affinity_probability_binary2"]
+                    affinity_summary["affinity_pred_value1"] = pred_affinity_value1.item()
+                    affinity_summary["affinity_probability_binary1"] = (
+                        pred_affinity_probability1.item()
+                    )
+                    affinity_summary["affinity_pred_value2"] = pred_affinity_value2.item()
+                    affinity_summary["affinity_probability_binary2"] = (
+                        pred_affinity_probability2.item()
+                    )
+
+                # Save the affinity summary in the ligand's directory
+                path = ligand_dir / f"affinity_{record.id}.json"
+                with path.open("w") as f:
+                    f.write(json.dumps(affinity_summary, indent=4))
+        else:
+            # Single ligand case, save in record directory
+            # Dump affinity summary
+            affinity_summary = {}
+            pred_affinity_value = prediction["affinity_pred_value"]
+            pred_affinity_probability = prediction["affinity_probability_binary"]
+            affinity_summary = {
+                "affinity_pred_value": pred_affinity_value.item(),
+                "affinity_probability_binary": pred_affinity_probability.item(),
+            }
+            if "affinity_pred_value1" in prediction:
+                pred_affinity_value1 = prediction["affinity_pred_value1"]
+                pred_affinity_probability1 = prediction["affinity_probability_binary1"]
+                pred_affinity_value2 = prediction["affinity_pred_value2"]
+                pred_affinity_probability2 = prediction["affinity_probability_binary2"]
+                affinity_summary["affinity_pred_value1"] = pred_affinity_value1.item()
+                affinity_summary["affinity_probability_binary1"] = (
+                    pred_affinity_probability1.item()
+                )
+                affinity_summary["affinity_pred_value2"] = pred_affinity_value2.item()
+                affinity_summary["affinity_probability_binary2"] = (
+                    pred_affinity_probability2.item()
+                )
+
+            # Save the affinity summary
+            record_dir.mkdir(exist_ok=True)
+            path = record_dir / f"affinity_{record.id}.json"
+            with path.open("w") as f:
+                f.write(json.dumps(affinity_summary, indent=4))
 
     def on_predict_epoch_end(
         self,

{boltz_vsynthes-1.0.34 → boltz_vsynthes-1.0.35/src/boltz_vsynthes.egg-info}/PKG-INFO

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