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

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

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

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

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

{boltz_vsynthes-1.0.35 → boltz_vsynthes-1.0.37}/src/boltz/data/parse/schema.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 import json
+import yaml
 
 import click
 import numpy as np
@@ -936,6 +937,7 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
     ccd: Mapping[str, Mol],
     mol_dir: Optional[Path] = None,
     boltz_2: bool = False,
+    output_dir: Optional[Path] = None,
 ) -> Target:
     """Parse a Boltz input yaml / json.
 
@@ -987,6 +989,8 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
         Path to the directory containing the molecules.
     boltz2: bool
         Whether to parse the input for Boltz2.
+    output_dir: Path, optional
+        Path to the output directory. If provided, results will be saved in a subfolder named after the input file.
 
     Returns
     -------
@@ -994,6 +998,14 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
         The parsed target.
 
     """
+    # Create output directory if specified
+    if output_dir is not None:
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        # Create subfolder based on input name
+        subfolder = output_dir / name
+        subfolder.mkdir(parents=True, exist_ok=True)
+
     # Assert version 1
     version = schema.get("version", 1)
     if version != 1:
@@ -1222,7 +1234,7 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
             )
 
         # Parse a non-polymer
-        elif (entity_type == "ligand") and ("ccd" in items[0][entity_type]):
+        elif (entity_type == "ligand") and "ccd" in (items[0][entity_type]):
             seq = items[0][entity_type]["ccd"]
 
             if isinstance(seq, str):
@@ -1314,60 +1326,6 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
                 "cyclic", False
             ), "Cyclic flag is not supported for ligands"
 
-        elif (entity_type == "ligand") and ("sdf" in items[0][entity_type]):
-            # Handle SDF file
-            sdf_path = Path(items[0][entity_type]["sdf"])
-            from boltz.data.parse.sdf import parse_sdf
-            target = parse_sdf(sdf_path, ccd, mol_dir)
-            mol = target["sequences"][0]["ligand"]["smiles"]
-
-            if affinity:
-                mol = standardize(mol)
-
-            mol = AllChem.MolFromSmiles(mol)
-            mol = AllChem.AddHs(mol)
-
-            # Set atom names
-            canonical_order = AllChem.CanonicalRankAtoms(mol)
-            for atom, can_idx in zip(mol.GetAtoms(), canonical_order):
-                atom_name = atom.GetSymbol().upper() + str(can_idx + 1)
-                if len(atom_name) > 4:
-                    msg = (
-                        f"{mol} has an atom with a name longer than "
-                        f"4 characters: {atom_name}."
-                    )
-                    raise ValueError(msg)
-                atom.SetProp("name", atom_name)
-
-            success = compute_3d_conformer(mol)
-            if not success:
-                msg = f"Failed to compute 3D conformer for {mol}"
-                raise ValueError(msg)
-
-            mol_no_h = AllChem.RemoveHs(mol, sanitize=False)
-            affinity_mw = AllChem.Descriptors.MolWt(mol_no_h) if affinity else None
-            extra_mols[f"LIG{ligand_id}"] = mol_no_h
-            residue = parse_ccd_residue(
-                name=f"LIG{ligand_id}",
-                ref_mol=mol,
-                res_idx=0,
-            )
-
-            ligand_id += 1
-            parsed_chain = ParsedChain(
-                entity=entity_id,
-                residues=[residue],
-                type=const.chain_type_ids["NONPOLYMER"],
-                cyclic_period=0,
-                sequence=None,
-                affinity=affinity,
-                affinity_mw=affinity_mw,
-            )
-
-            assert not items[0][entity_type].get(
-                "cyclic", False
-            ), "Cyclic flag is not supported for ligands"
-
         else:
             msg = f"Invalid entity type: {entity_type}"
             raise ValueError(msg)
@@ -1426,43 +1384,24 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
             protein_chains.add(chain_name)
 
         # Add affinity info
+        if chain.affinity and affinity_info is not None:
+            msg = "Cannot compute affinity for multiple ligands!"
+            raise ValueError(msg)
+
         if chain.affinity:
-            # If this is a protein binder, we need to create affinity info for each ligand
-            if chain_name in affinity_proteins:
-                # Find all ligand chains
-                ligand_chains = [
-                    (name, c) for name, c in chains.items() 
-                    if c.type == const.chain_type_ids["NONPOLYMER"]
-                ]
-                if not ligand_chains:
-                    msg = "No ligand chains found for protein binder!"
-                    raise ValueError(msg)
-                
-                # Create affinity info for each ligand
-                for ligand_name, ligand_chain in ligand_chains:
-                    affinity_info = AffinityInfo(
-                        chain_id=asym_id,
-                        mw=chain.affinity_mw,
-                    )
-                    # Save the affinity info in a subfolder named after the ligand
-                    output_dir = Path(f"output/{ligand_name}")
-                    output_dir.mkdir(parents=True, exist_ok=True)
-                    # Save the affinity info
-                    with open(output_dir / "affinity_info.json", "w") as f:
-                        json.dump({
-                            "chain_id": asym_id,
-                            "mw": chain.affinity_mw,
-                            "ligand_name": ligand_name
-                        }, f)
-            else:
-                # This is a ligand binder
-                if affinity_info is not None:
-                    msg = "Cannot compute affinity for multiple ligands!"
-                    raise ValueError(msg)
-                affinity_info = AffinityInfo(
-                    chain_id=asym_id,
-                    mw=chain.affinity_mw,
-                )
+            affinity_info = AffinityInfo(
+                chain_id=asym_id,
+                mw=chain.affinity_mw,
+            )
+            # Save affinity info if output directory is specified
+            if output_dir is not None:
+                affinity_path = subfolder / "affinity_info.json"
+                with open(affinity_path, "w") as f:
+                    json.dump({
+                        "chain_id": asym_id,
+                        "mw": chain.affinity_mw,
+                        "chain_name": chain_name
+                    }, f)
 
         # Find all copies of this chain in the assembly
         entity_id = int(chain.entity)
@@ -1925,3 +1864,64 @@ def standardize(smiles: str) -> Optional[str]:
         raise ValueError("Molecule is broken")
 
     return smiles
+
+
+def parse_boltz_directory(
+    input_dir: Path,
+    output_dir: Path,
+    ccd: Mapping[str, Mol],
+    mol_dir: Optional[Path] = None,
+    boltz_2: bool = False,
+) -> list[Target]:
+    """Parse all YAML files in a directory.
+
+    Parameters
+    ----------
+    input_dir : Path
+        Path to the directory containing YAML files.
+    output_dir : Path
+        Path to the output directory where results will be saved.
+    ccd : Mapping[str, Mol]
+        Dictionary of CCD components.
+    mol_dir : Path, optional
+        Path to the directory containing the molecules.
+    boltz_2 : bool, optional
+        Whether to parse the input for Boltz2.
+
+    Returns
+    -------
+    list[Target]
+        List of parsed targets.
+
+    """
+    input_dir = Path(input_dir)
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    targets = []
+    for yaml_file in input_dir.glob("*.yaml"):
+        # Skip hidden files and directories
+        if yaml_file.name.startswith('.') or any(part.startswith('.') for part in yaml_file.parts):
+            continue
+            
+        try:
+            # Load YAML file
+            with open(yaml_file, "r") as f:
+                schema = yaml.safe_load(f)
+
+            # Parse schema
+            target = parse_boltz_schema(
+                name=yaml_file.stem,
+                schema=schema,
+                ccd=ccd,
+                mol_dir=mol_dir,
+                boltz_2=boltz_2,
+                output_dir=output_dir,
+            )
+            targets.append(target)
+
+        except Exception as e:
+            print(f"Error processing {yaml_file}: {str(e)}")
+            continue
+
+    return targets

{boltz_vsynthes-1.0.35 → boltz_vsynthes-1.0.37}/src/boltz/data/parse/yaml.py

@@ -1,9 +1,10 @@
 from pathlib import Path
+from typing import Union, List, Optional
 
 import yaml
 from rdkit.Chem.rdchem import Mol
 
-from boltz.data.parse.schema import parse_boltz_schema
+from boltz.data.parse.schema import parse_boltz_schema, parse_boltz_directory
 from boltz.data.types import Target
 
 
@@ -12,8 +13,9 @@ def parse_yaml(
     ccd: dict[str, Mol],
     mol_dir: Path,
     boltz2: bool = False,
-) -> Target:
-    """Parse a Boltz input yaml / json.
+    output_dir: Optional[Path] = None,
+) -> Union[Target, List[Target]]:
+    """Parse a Boltz input yaml / json file or directory.
 
     The input file should be a yaml file with the following format:
 
@@ -49,20 +51,28 @@ def parse_yaml(
     Parameters
     ----------
     path : Path
-        Path to the YAML input format.
-    components : Dict
+        Path to the YAML input file or directory containing YAML files.
+    ccd : Dict
         Dictionary of CCD components.
-    boltz2 : bool
+    mol_dir : Path
+        Path to the directory containing molecules.
+    boltz2 : bool, optional
         Whether to parse the input for Boltz2.
+    output_dir : Path, optional
+        Path to the output directory where results will be saved.
 
     Returns
     -------
-    Target
-        The parsed target.
+    Union[Target, List[Target]]
+        The parsed target(s).
 
     """
-    with path.open("r") as file:
-        data = yaml.safe_load(file)
-
-    name = path.stem
-    return parse_boltz_schema(name, data, ccd, mol_dir, boltz2)
+    path = Path(path)
+    
+    if path.is_dir():
+        return parse_boltz_directory(path, output_dir or path, ccd, mol_dir, boltz2)
+    else:
+        with path.open("r") as file:
+            data = yaml.safe_load(file)
+        name = path.stem
+        return parse_boltz_schema(name, data, ccd, mol_dir, boltz2, output_dir)

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

@@ -290,72 +290,35 @@ class BoltzAffinityWriter(BasePredictionWriter):
         if prediction["exception"]:
             self.failed += 1
             return
-
-        # 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))
+        # 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))
 
     def on_predict_epoch_end(
         self,

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

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