Entities Module

This module contains entity-related functionality for handling ligands and molecular structures.

Ligand Utilities

ligand

Ligand processing and CCD cache management for Boltz predictions.

This module provides utilities for processing molecular structures, generating conformers, and managing the Chemical Component Dictionary (CCD) cache used by Boltz for ligand predictions.

sanitize_mol_id

sanitize_mol_id(mol_id: str) -> str

Ensure molecule ID complies with CCD naming rules (max 5 characters).

Parameters:

Name	Type	Description	Default
mol_id	str	Original molecule identifier.	required

Returns:

Type	Description
str	Sanitized molecule ID, truncated to 5 characters if necessary.

Warnings

Logs a warning if the ID is truncated.

Source code in src/cofolder/modules/entities/ligand.py

def sanitize_mol_id(mol_id: str) -> str:
    """Ensure molecule ID complies with CCD naming rules (max 5 characters).

    Parameters
    ----------
    mol_id : str
        Original molecule identifier.

    Returns
    -------
    str
        Sanitized molecule ID, truncated to 5 characters if necessary.

    Warnings
    --------
    Logs a warning if the ID is truncated.
    """
    if len(mol_id) > 5:
        truncated = mol_id[:5]
        logger.warning(f"[WARNING] Molecule ID '{mol_id}' is longer than 5 characters. "
              f"Truncating to '{truncated}' to comply with CCD naming rules.")
        return truncated
    return mol_id

add_pickled_prop

add_pickled_prop(mol: Mol, name: str, value)

Pickle and attach a property to an RDKit molecule.

Serializes a value using pickle and stores it as a hex-encoded string property on the molecule. Used for storing complex data structures in CCD format.

Parameters:

Name	Type	Description	Default
mol	Mol	RDKit molecule object to modify.	required
name	str	Property name to use for storage.	required
value	any	Value to pickle and attach (can be any picklable Python object).	required

Source code in src/cofolder/modules/entities/ligand.py

def add_pickled_prop(mol: Chem.Mol, name: str, value):
    """Pickle and attach a property to an RDKit molecule.

    Serializes a value using pickle and stores it as a hex-encoded
    string property on the molecule. Used for storing complex data
    structures in CCD format.

    Parameters
    ----------
    mol : Chem.Mol
        RDKit molecule object to modify.
    name : str
        Property name to use for storage.
    value : any
        Value to pickle and attach (can be any picklable Python object).
    """
    mol.SetProp(name, pickle.dumps(value).hex())

extract_constraints

extract_constraints(constraints, key, transpose=False)

Extract a list of values from constraint objects.

Parameters:

Name	Type	Description	Default
constraints	list	List of constraint objects.	required
key	str	Attribute name to extract.	required
transpose	bool	Whether to transpose the result (default False).	False

Returns:

Type	Description
list or ndarray	Extracted values, optionally transposed.

Source code in src/cofolder/modules/entities/ligand.py

def extract_constraints(constraints, key, transpose=False):
    """
    Extract a list of values from constraint objects.

    Parameters
    ----------
    constraints : list
        List of constraint objects.
    key : str
        Attribute name to extract.
    transpose : bool, optional
        Whether to transpose the result (default False).

    Returns
    -------
    list or numpy.ndarray
        Extracted values, optionally transposed.
    """
    result = [getattr(obj, key) for obj in constraints]
    if transpose:
        import numpy as np
        return np.array(result).transpose()
    return result

generate_2d_conformers

generate_2d_conformers(sdf_in: str, sdf_out: Optional[str] = None)

Generate 2D coordinate layouts for molecules in an SDF file.

Computes 2D coordinates suitable for visualization and quick analysis. Useful for preparing ligands that only have SMILES or 1D representations.

Parameters:

Name	Type	Description	Default
sdf_in	str	Path to input SDF file.	required
sdf_out	str	Path to output SDF file. If None, overwrites the input file.	None

Notes

Uses RDKit's Compute2DCoords for layout generation. Invalid molecules (None) are skipped.

Source code in src/cofolder/modules/entities/ligand.py

def generate_2d_conformers(sdf_in: str, sdf_out: Optional[str] = None):
    """Generate 2D coordinate layouts for molecules in an SDF file.

    Computes 2D coordinates suitable for visualization and quick analysis.
    Useful for preparing ligands that only have SMILES or 1D representations.

    Parameters
    ----------
    sdf_in : str
        Path to input SDF file.
    sdf_out : str, optional
        Path to output SDF file. If None, overwrites the input file.

    Notes
    -----
    Uses RDKit's Compute2DCoords for layout generation.
    Invalid molecules (None) are skipped.
    """
    sdf_out = sdf_out or sdf_in
    mols = read.read_sdf(sdf_in)

    for mol in mols:
        rdDepictor.Compute2DCoords(mol)

    write.write_sdf(mols, sdf_out)
    logger.info(f"Generated 2D conformers in {sdf_out}")

generate_3d_conformers

generate_3d_conformers(sdf_in: str, sdf_out: Optional[str] = None)

Generate 3D conformers for molecules using ETKDG and UFF optimization.

Creates physically realistic 3D conformations by: 1. Adding hydrogens 2. Embedding with ETKDG (Experimental-Torsion Knowledge Distance Geometry) 3. Optimizing geometry with UFF (Universal Force Field)

Parameters:

Name	Type	Description	Default
sdf_in	str	Path to input SDF file.	required
sdf_out	str	Path to output SDF file. If None, overwrites the input file.	None

Warnings

Logs a warning if 3D conformer generation fails for any molecule. Failed molecules are skipped and not written to output.

Notes

This is computationally more expensive than 2D conformer generation but produces geometries suitable for docking and structure prediction.

Source code in src/cofolder/modules/entities/ligand.py

def generate_3d_conformers(sdf_in: str, sdf_out: Optional[str] = None):
    """Generate 3D conformers for molecules using ETKDG and UFF optimization.

    Creates physically realistic 3D conformations by:
    1. Adding hydrogens
    2. Embedding with ETKDG (Experimental-Torsion Knowledge Distance Geometry)
    3. Optimizing geometry with UFF (Universal Force Field)

    Parameters
    ----------
    sdf_in : str
        Path to input SDF file.
    sdf_out : str, optional
        Path to output SDF file. If None, overwrites the input file.

    Warnings
    --------
    Logs a warning if 3D conformer generation fails for any molecule.
    Failed molecules are skipped and not written to output.

    Notes
    -----
    This is computationally more expensive than 2D conformer generation
    but produces geometries suitable for docking and structure prediction.
    """
    sdf_out = sdf_out or sdf_in
    mols = read.read_sdf(sdf_in)
    mols_3d = []

    for mol in mols:
        try:
            mol = Chem.AddHs(mol)
            AllChem.EmbedMolecule(mol, AllChem.ETKDG())
            AllChem.UFFOptimizeMolecule(mol)
            mols_3d.append(mol)
        except Exception as e:
            logger.warning(f"3D conformer generation failed for molecule: {e}")

    write.write_sdf(mols_3d, sdf_out)
    logger.info(f"Generated 3D conformers in {sdf_out}")

iterate_sdf_records

iterate_sdf_records(sdf_path: str, id_property: str)

Compatibility iterator that now retains record indexes while yielding valid molblocks.

Source code in src/cofolder/modules/entities/ligand.py

def iterate_sdf_records(sdf_path: str, id_property: str):
    """Compatibility iterator that now retains record indexes while yielding valid molblocks."""
    for record in iter_sdf_records(sdf_path):
        source_id = record.title if id_property == "_Name" else record.properties.get(id_property)
        mol_id = source_id or f"mol_{record.index}"
        try:
            mol = Chem.MolFromMolBlock(record.molblock, removeHs=False)
        except Exception:
            mol = None
        if mol is not None:
            yield record.index, mol_id, Chem.MolToMolBlock(mol)

csv_to_sdf

csv_to_sdf(
    csv_path: str,
    smiles_col: str,
    output_sdf_path: str | None = None,
    property_cols: Optional[Union[dict, List[dict], str, List[str]]] = None,
) -> None

Compatibility wrapper that preserves the older CSV-specific API.

Source code in src/cofolder/modules/entities/ligand.py

def csv_to_sdf(
    csv_path: str,
    smiles_col: str,
    output_sdf_path: str | None = None,
    property_cols: Optional[Union[dict, List[dict], str, List[str]]] = None,
) -> None:
    """Compatibility wrapper that preserves the older CSV-specific API."""
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"CSV file not found: {csv_path}")

    return smiles_to_sdf(
        data=csv_path,
        smiles_col=smiles_col,
        output_sdf_path=output_sdf_path,
        property_cols=property_cols,
    )

mol_to_ccd

mol_to_ccd(
    resname: str,
    mol: Mol,
    boltz_path: Union[str, PathLike] = os.path.expanduser("~/.boltz"),
)

Convert an RDKit molecule to Boltz CCD format and cache it.

Processes an RDKit molecule by: 1. Sanitizing and preparing the structure 2. Parsing geometric and stereochemical constraints 3. Pickling all properties for Boltz compatibility 4. Saving to the CCD cache directory

Parameters:

Name	Type	Description	Default
resname	str	Residue name (identifier) for the molecule. Must be 5 characters or less.	required
mol	Mol	RDKit molecule object to convert.	required
boltz_path	str or Path	Path to Boltz cache directory where CCD files are stored.	"~/.boltz"

Raises:

Type	Description
Exception	If the molecule is None.

Examples:

>>> from rdkit import Chem
>>> from cofolder.modules.entities.ligand import mol_to_ccd
>>> smiles = "CCO"
>>> mol = Chem.MolFromSmiles(smiles)
>>> mol_to_ccd("ETH", mol)
# Saves to ~/.boltz/mols/ETH.pkl

Notes

This function contains adapted logic from jacktday/boltztools (MIT).

The CCD format includes: - Atomic symmetries - Bond length and angle constraints - Chiral center definitions - Stereochemical bond configurations - Aromatic ring planarity constraints

Source code in src/cofolder/modules/entities/ligand.py

def mol_to_ccd(resname: str, mol: Chem.Mol, boltz_path: Union[str, os.PathLike] = os.path.expanduser("~/.boltz")):
    """Convert an RDKit molecule to Boltz CCD format and cache it.

    Processes an RDKit molecule by:
    1. Sanitizing and preparing the structure
    2. Parsing geometric and stereochemical constraints
    3. Pickling all properties for Boltz compatibility
    4. Saving to the CCD cache directory

    Parameters
    ----------
    resname : str
        Residue name (identifier) for the molecule. Must be 5 characters or less.
    mol : Chem.Mol
        RDKit molecule object to convert.
    boltz_path : str or Path, default="~/.boltz"
        Path to Boltz cache directory where CCD files are stored.

    Raises
    ------
    Exception
        If the molecule is None.

    Examples
    --------
    >>> from rdkit import Chem
    >>> from cofolder.modules.entities.ligand import mol_to_ccd
    >>> smiles = "CCO"
    >>> mol = Chem.MolFromSmiles(smiles)
    >>> mol_to_ccd("ETH", mol)
    # Saves to ~/.boltz/mols/ETH.pkl

    Notes
    -----
    This function contains adapted logic from jacktday/boltztools (MIT).

    The CCD format includes:
    - Atomic symmetries
    - Bond length and angle constraints
    - Chiral center definitions
    - Stereochemical bond configurations
    - Aromatic ring planarity constraints
    """
    if not isinstance(resname, str) or not resname or len(resname) > 5 or not resname.isalnum():
        raise ValueError(
            "CCD identifier must be a nonempty alphanumeric code of at most 5 characters."
        )
    if mol is None:
        raise Exception("Mol can not be null")

    mol = Chem.RemoveHs(mol)
    mol.UpdatePropertyCache(strict=False)
    Chem.SanitizeMol(mol)
    rdmolops.AssignStereochemistryFrom3D(mol)

    mol = _prepare_mol(mol)

    parse_ccd_residue = _load_parse_ccd_residue()
    parsedResidue = parse_ccd_residue(resname, mol, 0)

    add_pickled_prop(mol, 'MOL_NAME', resname)
    add_pickled_prop(mol, 'symmetries', mol.GetSubstructMatches(mol, uniquify=False))
    add_pickled_prop(mol, 'pb_edge_index', extract_constraints(parsedResidue.rdkit_bounds_constraints, 'atom_idxs', transpose=True))
    add_pickled_prop(mol, 'pb_lower_bounds', extract_constraints(parsedResidue.rdkit_bounds_constraints, 'lower_bound'))
    add_pickled_prop(mol, 'pb_upper_bounds', extract_constraints(parsedResidue.rdkit_bounds_constraints, 'upper_bound'))
    add_pickled_prop(mol, 'pb_bond_mask', extract_constraints(parsedResidue.rdkit_bounds_constraints, 'is_bond'))
    add_pickled_prop(mol, 'pb_angle_mask', extract_constraints(parsedResidue.rdkit_bounds_constraints, 'is_angle'))
    add_pickled_prop(mol, 'chiral_atom_index', extract_constraints(parsedResidue.chiral_atom_constraints, 'atom_idxs', transpose=True))
    add_pickled_prop(mol, 'chiral_check_mask', extract_constraints(parsedResidue.chiral_atom_constraints, 'is_reference'))
    add_pickled_prop(mol, 'chiral_atom_orientations', extract_constraints(parsedResidue.chiral_atom_constraints, 'is_r'))
    add_pickled_prop(mol, 'stereo_bond_index', extract_constraints(parsedResidue.stereo_bond_constraints, 'atom_idxs', transpose=True))
    add_pickled_prop(mol, 'stereo_check_mask', extract_constraints(parsedResidue.stereo_bond_constraints, 'is_check'))
    add_pickled_prop(mol, 'stereo_bond_orientations', extract_constraints(parsedResidue.stereo_bond_constraints, 'is_e'))
    add_pickled_prop(mol, 'aromatic_5_ring_index', extract_constraints(parsedResidue.planar_ring_5_constraints, 'atom_idxs', transpose=True))
    add_pickled_prop(mol, 'aromatic_6_ring_index', extract_constraints(parsedResidue.planar_ring_6_constraints, 'atom_idxs', transpose=True))
    add_pickled_prop(mol, 'planar_double_bond_index', extract_constraints(parsedResidue.planar_bond_constraints, 'atom_idxs', transpose=True))

    Chem.SetDefaultPickleProperties(Chem.PropertyPickleOptions.AllProps)
    mols_dir = Path(boltz_path) / 'mols'
    mols_dir.mkdir(parents=True, exist_ok=True)

    logger.debug("Checking atom properties just before pickle:")
    atoms = list(mol.GetAtoms())
    logger.debug("Total atoms: %d", len(atoms))
    missing = []
    for atom in atoms:
        if not atom.HasProp("name"):
            missing.append(atom.GetIdx())
        # show first few atoms for inspection
        logger.debug(f"  idx {atom.GetIdx()} props: {atom.GetPropsAsDict()}")

    if missing:
        logger.warning(f"{len(missing)} atoms are missing 'name' property. Example indices: {missing[:10]}")
    else:
        logger.debug("All atoms have 'name' property.")

    write.write_pickle(mol, mols_dir / f"{resname}.pkl")