boltz-vsynthes 1.0.34__py3-none-any.whl → 1.0.35__py3-none-any.whl

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.dist-info → boltz_vsynthes-1.0.35.dist-info}/METADATA

@@ -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.dist-info → boltz_vsynthes-1.0.35.dist-info}/RECORD

@@ -56,7 +56,7 @@ boltz/data/write/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 boltz/data/write/mmcif.py,sha256=zHwe4ulietjCVz5JMzl4_bu4liqwwcUKlgr9orVVd2Q,12177
 boltz/data/write/pdb.py,sha256=P2pDpKKIJuIoPADWwei0tIakHWE-ajlGEw6BE371EcI,6235
 boltz/data/write/utils.py,sha256=2PzveM4qZ7Aadru1Qzoi2Dhit_s4LXlz4ceXpGJi4y0,539
-boltz/data/write/writer.py,sha256=bYxHHf6e0waO8WdAP_AmzWEEljc8lRx85HcDPR2Tjqo,12500
+boltz/data/write/writer.py,sha256=Q6h6R1pU-wWJmdBnAsxjOpDwrVKjlfFKy8rTbL9_zfE,14644
 boltz/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/model/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/model/layers/attention.py,sha256=rkd5PTIoVB1kSzDnBrZfut4DoEHKHBCwuDb1p0JfSck,3782
@@ -107,9 +107,9 @@ boltz/model/optim/scheduler.py,sha256=nB4jz0CZ4pR4n08LQngExL_pNycIdYI8AXVoHPnZWQ
 boltz/model/potentials/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 boltz/model/potentials/potentials.py,sha256=vev8Vjfs-ML1hyrdv_R8DynG4wSFahJ6nzPWp7CYQqw,17507
 boltz/model/potentials/schedules.py,sha256=m7XJjfuF9uTX3bR9VisXv1rvzJjxiD8PobXRpcBBu1c,968
-boltz_vsynthes-1.0.34.dist-info/licenses/LICENSE,sha256=8GZ_1eZsUeG6jdqgJJxtciWzADfgLEV4LY8sKUOsJhc,1102
-boltz_vsynthes-1.0.34.dist-info/METADATA,sha256=h-T8W9FPQG5hynUMxYWhnJN4Ztt1XWoX4LgorJCEEYY,7171
-boltz_vsynthes-1.0.34.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-boltz_vsynthes-1.0.34.dist-info/entry_points.txt,sha256=n5a5I35ntu9lmyr16oZgHPFY0b0YxjiixY7m7nbMTLc,41
-boltz_vsynthes-1.0.34.dist-info/top_level.txt,sha256=MgU3Jfb-ctWm07YGMts68PMjSh9v26D0gfG3dFRmVFA,6
-boltz_vsynthes-1.0.34.dist-info/RECORD,,
+boltz_vsynthes-1.0.35.dist-info/licenses/LICENSE,sha256=8GZ_1eZsUeG6jdqgJJxtciWzADfgLEV4LY8sKUOsJhc,1102
+boltz_vsynthes-1.0.35.dist-info/METADATA,sha256=XX9hHMNRhqMjmfhzZxJ0YtFYR9_ppHPra433DXQBxo0,7171
+boltz_vsynthes-1.0.35.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+boltz_vsynthes-1.0.35.dist-info/entry_points.txt,sha256=n5a5I35ntu9lmyr16oZgHPFY0b0YxjiixY7m7nbMTLc,41
+boltz_vsynthes-1.0.35.dist-info/top_level.txt,sha256=MgU3Jfb-ctWm07YGMts68PMjSh9v26D0gfG3dFRmVFA,6
+boltz_vsynthes-1.0.35.dist-info/RECORD,,