boltz-vsynthes 1.0.6__tar.gz → 1.0.8__tar.gz

{boltz_vsynthes-1.0.6/src/boltz_vsynthes.egg-info → boltz_vsynthes-1.0.8}/PKG-INFO

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

{boltz_vsynthes-1.0.6 → boltz_vsynthes-1.0.8}/pyproject.toml

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

{boltz_vsynthes-1.0.6 → boltz_vsynthes-1.0.8}/src/boltz/data/mol.py

@@ -30,6 +30,10 @@ def load_molecules(moldir: str, molecules: list[str]) -> dict[str, Mol]:
     """
     loaded_mols = {}
     for molecule in molecules:
+        # Skip if it's a SMILES string (starts with LIG)
+        if molecule.startswith("LIG"):
+            continue
+            
         path = Path(moldir) / f"{molecule}.pkl"
         if not path.exists():
             msg = f"CCD component {molecule} not found!"

{boltz_vsynthes-1.0.6 → boltz_vsynthes-1.0.8}/src/boltz/data/parse/schema.py

@@ -621,6 +621,9 @@ def get_mol(ccd: str, mols: dict, moldir: str) -> Mol:
     Return mol with ccd from mols if it is in mols. Otherwise load it from moldir,
     add it to mols, and return the mol.
     """
+    # Skip if it's a SMILES string (starts with LIG)
+    if ccd.startswith("LIG"):
+        return None
     mol = mols.get(ccd)
     if mol is None:
         mol = load_molecules(moldir, [ccd])[ccd]
@@ -655,6 +658,10 @@ def parse_ccd_residue(
        The output ParsedResidue, if successful.
 
     """
+    # Skip if it's a SMILES string (starts with LIG)
+    if name.startswith("LIG"):
+        return None
+        
     unk_chirality = const.chirality_type_ids[const.unk_chirality_type]
 
     # Check if this is a single heavy atom CCD residue
@@ -929,86 +936,32 @@ def token_spec_to_ids(
         contacts.append((chain_to_idx[chain_name], residue_index_or_atom_name - 1))
 
 
-def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
-    name: str,
-    schema: dict,
-    ccd: Mapping[str, Mol],
-    mol_dir: Optional[Path] = None,
-    boltz_2: bool = False,
-) -> Target:
-    """Parse a Boltz input yaml / json.
-
-    The input file should be a dictionary with the following format:
-
-    version: 1
-    sequences:
-        - protein:
-            id: A
-            sequence: "MADQLTEEQIAEFKEAFSLF"  # or pdb: "1a2k" or pdb: "path/to/file.pdb"
-            msa: path/to/msa1.a3m
-        - protein:
-            id: [B, C]
-            sequence: "AKLSILPWGHC"
-            msa: path/to/msa2.a3m
-        - rna:
-            id: D
-            sequence: "GCAUAGC"
-        - ligand:
-            id: E
-            smiles: "CC1=CC=CC=C1"
-    constraints:
-        - bond:
-            atom1: [A, 1, CA]
-            atom2: [A, 2, N]
-        - pocket:
-            binder: E
-            contacts: [[B, 1], [B, 2]]
-            max_distance: 6
-        - contact:
-            token1: [A, 1]
-            token2: [B, 1]
-            max_distance: 6
-    templates:
-        - cif: path/to/template.cif
-    properties:
-        - affinity:
-            binder: E
+def parse_boltz_schema(schema: dict) -> dict:
+    """Parse the Boltz input schema.
 
     Parameters
     ----------
-    name : str
-        A name for the input.
     schema : dict
         The input schema.
-    components : dict
-        Dictionary of CCD components.
-    mol_dir: Path
-        Path to the directory containing the molecules.
-    boltz2: bool
-        Whether to parse the input for Boltz2.
 
     Returns
     -------
-    Target
-        The parsed target.
+    dict
+        The parsed schema.
 
     """
-    # Assert version 1
-    version = schema.get("version", 1)
-    if version != 1:
-        msg = f"Invalid version {version} in input!"
+    # Check version
+    if "version" not in schema:
+        msg = "Schema must have a version field"
         raise ValueError(msg)
 
-    # Disable rdkit warnings
-    blocker = rdBase.BlockLogs()  # noqa: F841
-
-    # First group items that have the same type, sequence and modifications
+    # Group items by entity type and sequence
     items_to_group = {}
     chain_name_to_entity_type = {}
-    
+
     # Keep track of ligand IDs
     ligand_id = 1
-    ligand_id_map = {}  # Maps user-provided IDs to internal LIG1, LIG2, etc.
+    ligand_id_map = {}
 
     # Parse sequences
     for item in schema["sequences"]:
@@ -1057,16 +1010,37 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
                 msg = "Protein must have either 'sequence' or 'pdb' field"
                 raise ValueError(msg)
         elif entity_type == "ligand":
-            assert "smiles" in item[entity_type] or "ccd" in item[entity_type]
-            assert "smiles" not in item[entity_type] or "ccd" not in item[entity_type]
+            # Support for SMILES, CCD, and SDF
             if "smiles" in item[entity_type]:
                 seq = str(item[entity_type]["smiles"])
                 # Map user-provided ID to internal LIG1, LIG2, etc.
                 for id in entity_id:
                     ligand_id_map[id] = f"LIG{ligand_id}"
                 ligand_id += 1
-            else:
+            elif "ccd" in item[entity_type]:
                 seq = str(item[entity_type]["ccd"])
+                # For CCD ligands, use the CCD code as the internal ID
+                for id in entity_id:
+                    ligand_id_map[id] = seq
+            elif "sdf" in item[entity_type]:
+                sdf_path = Path(item[entity_type]["sdf"])
+                if not sdf_path.exists():
+                    msg = f"SDF file not found: {sdf_path}"
+                    raise FileNotFoundError(msg)
+                # Read SDF and convert to SMILES
+                from rdkit import Chem
+                mol = Chem.SDMolSupplier(str(sdf_path))[0]
+                if mol is None:
+                    msg = f"Failed to read SDF file: {sdf_path}"
+                    raise ValueError(msg)
+                seq = Chem.MolToSmiles(mol)
+                # Map user-provided ID to internal LIG1, LIG2, etc.
+                for id in entity_id:
+                    ligand_id_map[id] = f"LIG{ligand_id}"
+                ligand_id += 1
+            else:
+                msg = "Ligand must have either 'smiles', 'ccd', or 'sdf' field"
+                raise ValueError(msg)
 
         # Group items by entity
         items_to_group.setdefault((entity_type, seq), []).append(item)
@@ -1077,208 +1051,60 @@ def parse_boltz_schema(  # noqa: C901, PLR0915, PLR0912
         for chain_name in chain_names:
             chain_name_to_entity_type[chain_name] = entity_type
 
-    # Check if any affinity ligand is present
-    affinity_ligands = set()
-    properties = schema.get("properties", [])
-    for prop in properties:
-        if "affinity" in prop:
-            binder = prop["affinity"]["binder"]
-            if binder in ligand_id_map:
-                binder = ligand_id_map[binder]  # Convert to internal LIG1, LIG2, etc.
-            affinity_ligands.add(binder)
-
-    # Parse each group
-    chains = []
-    extra_mols = {}
-    for (entity_type, seq), items in items_to_group.items():
-        # Get entity id
-        entity_id = items[0][entity_type]["id"]
+    # Get all proteins and ligands
+    proteins = []
+    ligands = []
+    for item in schema["sequences"]:
+        entity_type = list(item.keys())[0]
+        entity_id = item[entity_type]["id"]
         entity_id = [entity_id] if isinstance(entity_id, str) else entity_id
-
-        # Check if this entity has affinity
-        affinity = any(entity in affinity_ligands for entity in entity_id)
-
-        # Parse a protein
         if entity_type == "protein":
-            # Get MSA
-            msa = items[0][entity_type].get("msa")
-            if msa is not None:
-                msa = Path(msa)
-                if not msa.exists():
-                    msg = f"MSA file not found: {msa}"
-                    raise FileNotFoundError(msg)
-                with msa.open("r") as f:
-                    msa_data = f.read()
-            else:
-                msa_data = None
-
-            # Parse sequence
-            residues = []
-            for res_idx, code in enumerate(seq):
-                # Get mol
-                ref_mol = get_mol(code, ccd, mol_dir)
-
-                # Parse residue
-                residue = parse_ccd_residue(
-                    name=code,
-                    ref_mol=ref_mol,
-                    res_idx=res_idx,
-                )
-                residues.append(residue)
-
-            # Create protein chain
-            parsed_chain = ParsedChain(
-                entity=entity_id,
-                residues=residues,
-                type=const.chain_type_ids["PROTEIN"],
-                cyclic_period=0,
-                sequence=seq,
-                affinity=affinity,
-                affinity_mw=None,
-            )
-
-        # Parse a non-polymer
-        elif (entity_type == "ligand") and "ccd" in (items[0][entity_type]):
-            seq = items[0][entity_type]["ccd"]
-
-            if isinstance(seq, str):
-                seq = [seq]
-
-            if affinity and len(seq) > 1:
-                msg = "Cannot compute affinity for multi residue ligands!"
-                raise ValueError(msg)
-
-            residues = []
-            affinity_mw = None
-            for res_idx, code in enumerate(seq):
-                # Get mol
-                ref_mol = get_mol(code, ccd, mol_dir)
-
-                if affinity:
-                    affinity_mw = AllChem.Descriptors.MolWt(ref_mol)
-
-                # Parse residue
-                residue = parse_ccd_residue(
-                    name=code,
-                    ref_mol=ref_mol,
-                    res_idx=res_idx,
-                )
-                residues.append(residue)
-
-            # Create multi ligand chain
-            parsed_chain = ParsedChain(
-                entity=entity_id,
-                residues=residues,
-                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"
-
-        elif (entity_type == "ligand") and ("smiles" in items[0][entity_type]):
-            seq = items[0][entity_type]["smiles"]
-
-            if affinity:
-                seq = standardize(seq)
-
-            mol = AllChem.MolFromSmiles(seq)
-            mol = AllChem.AddHs(mol)
+            proteins.extend(entity_id)
+        elif entity_type == "ligand":
+            ligands.extend(entity_id)
 
-            # 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"{seq} has an atom with a name longer than "
-                        f"4 characters: {atom_name}."
-                    )
+    # Generate properties for each protein-ligand pair
+    new_properties = []
+    for prop in schema.get("properties", []):
+        if "affinity" in prop:
+            affinity = prop["affinity"]
+            # Handle protein as binder
+            if "protein" in affinity:
+                binder = affinity["protein"]
+                if binder not in proteins:
+                    msg = f"Protein {binder} not found in sequences"
                     raise ValueError(msg)
-                atom.SetProp("name", atom_name)
-
-            success = compute_3d_conformer(mol)
-            if not success:
-                msg = f"Failed to compute 3D conformer for {seq}"
-                raise ValueError(msg)
-
-            mol_no_h = AllChem.RemoveHs(mol, sanitize=False)
-            affinity_mw = AllChem.Descriptors.MolWt(mol_no_h) if affinity else None
-            
-            # Use the mapped internal ID (LIG1, LIG2, etc.)
-            internal_id = ligand_id_map[entity_id[0]]
-            extra_mols[internal_id] = mol_no_h
-            residue = parse_ccd_residue(
-                name=internal_id,
-                ref_mol=mol,
-                res_idx=0,
-            )
-
-            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)
-
-        chains.append(parsed_chain)
-
-    # Parse constraints
-    constraints = []
-    for constraint in schema.get("constraints", []):
-        if "bond" in constraint:
-            atom1 = constraint["bond"]["atom1"]
-            atom2 = constraint["bond"]["atom2"]
-            constraints.append(ParsedBond(atom1, atom2))
-        elif "pocket" in constraint:
-            binder = constraint["pocket"]["binder"]
-            if binder in ligand_id_map:
-                binder = ligand_id_map[binder]  # Convert to internal LIG1, LIG2, etc.
-            contacts = constraint["pocket"]["contacts"]
-            max_distance = constraint["pocket"].get("max_distance", 6.0)
-            constraints.append(ParsedPocket(binder, contacts, max_distance))
-        elif "contact" in constraint:
-            token1 = constraint["contact"]["token1"]
-            token2 = constraint["contact"]["token2"]
-            max_distance = constraint["contact"].get("max_distance", 6.0)
-            constraints.append(ParsedContact(token1, token2, max_distance))
-        else:
-            msg = f"Invalid constraint type: {list(constraint.keys())[0]}"
-            raise ValueError(msg)
-
-    # Parse templates
-    templates = []
-    for template in schema.get("templates", []):
-        cif = template["cif"]
-        chain_id = template.get("chain_id")
-        template_id = template.get("template_id")
-        templates.append(ParsedTemplate(cif, chain_id, template_id))
-
-    # Create target
-    target = Target(
-        name=name,
-        chains=chains,
-        constraints=constraints,
-        templates=templates,
-        extra_mols=extra_mols,
-    )
-
-    return target
+                # Generate pairs with all ligands
+                for ligand in ligands:
+                    if ligand in ligand_id_map:
+                        ligand = ligand_id_map[ligand]  # Convert to internal LIG1, LIG2, etc.
+                    new_properties.append({
+                        "affinity": {
+                            "binder": binder,
+                            "ligand": ligand
+                        }
+                    })
+            # Handle ligand as binder (backward compatibility)
+            elif "binder" in affinity:
+                binder = affinity["binder"]
+                if binder not in proteins:
+                    msg = f"Protein {binder} not found in sequences"
+                    raise ValueError(msg)
+                # Generate pairs with all ligands
+                for ligand in ligands:
+                    if ligand in ligand_id_map:
+                        ligand = ligand_id_map[ligand]  # Convert to internal LIG1, LIG2, etc.
+                    new_properties.append({
+                        "affinity": {
+                            "binder": binder,
+                            "ligand": ligand
+                        }
+                    })
+
+    # Update schema with generated properties
+    schema["properties"] = new_properties
+
+    return schema
 
 
 def standardize(smiles: str) -> Optional[str]:

{boltz_vsynthes-1.0.6 → boltz_vsynthes-1.0.8}/src/boltz/main.py

@@ -1235,72 +1235,89 @@ def predict(  # noqa: C901, PLR0915, PLR0912
         # 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)
-
-        pred_writer = BoltzAffinityWriter(
-            data_dir=processed.targets_dir,
-            output_dir=out_dir / "predictions",
-        )
+        # Group records by protein-ligand pairs
+        affinity_groups = {}
+        for record in manifest.records:
+            if record.affinity:
+                key = (record.affinity["binder"], record.affinity["ligand"])
+                if key not in affinity_groups:
+                    affinity_groups[key] = []
+                affinity_groups[key].append(record)
+
+        # Process each protein-ligand pair
+        for (binder, ligand), records in affinity_groups.items():
+            # Create subfolder for this protein-ligand pair
+            pair_dir = out_dir / "predictions" / f"{binder}_{ligand}"
+            pair_dir.mkdir(parents=True, exist_ok=True)
+
+            # Create manifest for this pair
+            pair_manifest = Manifest(records)
+
+            # Validate inputs
+            pair_manifest_filtered = filter_inputs_affinity(
+                manifest=pair_manifest,
+                outdir=pair_dir,
+                override=override,
+            )
+            if not pair_manifest_filtered.records:
+                click.echo(f"Found existing affinity predictions for {binder}_{ligand}, skipping.")
+                continue
 
-        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,
-        )
+            msg = f"Running affinity prediction for {binder} with {ligand}"
+            click.echo(msg)
 
-        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,
-        }
+            pred_writer = BoltzAffinityWriter(
+                data_dir=processed.targets_dir,
+                output_dir=pair_dir,
+            )
 
-        # Load affinity model
-        if affinity_checkpoint is None:
-            affinity_checkpoint = cache / "boltz2_aff.ckpt"
+            data_module = Boltz2InferenceDataModule(
+                manifest=pair_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-1.0.6 → boltz_vsynthes-1.0.8/src/boltz_vsynthes.egg-info}/PKG-INFO

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