Function multi_pose_sampling

flowdock/utils/sampling_utils.py:122–427

Sample multiple poses of a protein-ligand complex. :param receptor_path: Path to the receptor file. :param ligand_path: Path to the ligand file. :param cfg: Config dictionary. :param lit_module: LightningModule instance. :param out_path: Path to save the output files. :param

(
    receptor_path: str,
    ligand_path: str,
    cfg: DictConfig,
    lit_module: LightningModule,
    out_path: str,
    save_pdb: bool = True,
    separate_pdb: bool = True,
    chain_id: Optional[str] = None,
    apo_receptor_path: Optional[str] = None,
    sample_id: Optional[str] = None,
    protein: Optional[FDProtein] = None,
    sequences_to_embeddings: Optional[Dict[str, np.ndarray]] = None,
    confidence: bool = True,
    affinity: bool = True,
    return_all_states: bool = False,
    auxiliary_estimation_only: bool = False,
    **kwargs: Dict[str, Any],
)

def multi_pose_sampling(
    receptor_path: str,
    ligand_path: str,
    cfg: DictConfig,
    lit_module: LightningModule,
    out_path: str,
    save_pdb: bool = True,
    separate_pdb: bool = True,
    chain_id: Optional[str] = None,
    apo_receptor_path: Optional[str] = None,
    sample_id: Optional[str] = None,
    protein: Optional[FDProtein] = None,
    sequences_to_embeddings: Optional[Dict[str, np.ndarray]] = None,
    confidence: bool = True,
    affinity: bool = True,
    return_all_states: bool = False,
    auxiliary_estimation_only: bool = False,
    **kwargs: Dict[str, Any],
) -> Tuple[
    Optional[Chem.Mol],
    Optional[List[float]],
    Optional[List[float]],
    Optional[List[float]],
    Optional[List[Any]],
    Optional[Any],
    Optional[np.ndarray],
    Optional[np.ndarray],
]:
    """Sample multiple poses of a protein-ligand complex.

    :param receptor_path: Path to the receptor file.
    :param ligand_path: Path to the ligand file.
    :param cfg: Config dictionary.
    :param lit_module: LightningModule instance.
    :param out_path: Path to save the output files.
    :param save_pdb: Whether to save PDB files.
    :param separate_pdb: Whether to save separate PDB files for each pose.
    :param chain_id: Chain ID of the receptor.
    :param apo_receptor_path: Path to the optional apo receptor file.
    :param sample_id: Optional sample ID.
    :param protein: Optional protein object.
    :param sequences_to_embeddings: Mapping of sequences to embeddings.
    :param confidence: Whether to estimate confidence scores.
    :param affinity: Whether to estimate affinity scores.
    :param return_all_states: Whether to return all states.
    :param auxiliary_estimation_only: Whether to only estimate auxiliary outputs (e.g., confidence,
        affinity) for the input (generated) samples (potentially derived from external sources).
    :param kwargs: Additional keyword arguments.
    :return: Reference molecule, protein plDDTs, ligand plDDTs, ligand fragment plDDTs, estimated
        binding affinities, structure trajectories, input batch, B-factors, and structure rankings.
    """
    if return_all_states and auxiliary_estimation_only:
        # NOTE: If auxiliary estimation is solely enabled, structure trajectory sampling will be disabled
        return_all_states = False
    struct_res_all, lig_res_all = [], []
    plddt_all, plddt_lig_all, plddt_ligs_all, res_plddt_all = [], [], [], []
    affinity_all, ligs_affinity_all = [], []
    frames_all = []